def visited_base

bendtorch-core@0.1.0.0 · src/autograd/reduce_fiber_proof.bend:606

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/broadcast_proofs.bend as Broadcast
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

606def visited_base(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat) ->
607  {visited(rows, extents, fuel, target, base, digit) == shift(base, visited(rows, extents, fuel, target, 0n, digit)) : List<&2, Nat>}:
608  match rows extents fuel:
609    case Nil{} _ _:
610      %Equal.sym(Nat, Nat.add(base, 0n), base, Broadcast.add_zero(base)) : {[base] == [_] : List<&2, Nat>}
611      {==}
612    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
613      {==}
614    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
615      +inner = Nat.mod(target, left)
616      +offset = Nat.mul(right, digit)
617      +span = Reduce.first_span(rest, others)
618      +first = Reduce.first_digit(rest, inner)
619      +low = visited(rest, others, span, inner, 0n, first)
620      +high = visited(rows, extents, remaining, target, 0n, 1n+digit)
621      %Equal.sym(List<&2, Nat>, visited(rest, others, span, inner, Nat.add(base, offset), first), shift(Nat.add(base, offset), low),
622        visited_base(rest, others, span, inner, Nat.add(base, offset), first)) :
623        {List.append(&2, Nat, _, visited(rows, extents, remaining, target, base, 1n+digit)) ==
624          shift(base, List.append(&2, Nat, visited(rest, others, span, inner, offset, first), high)) : List<&2, Nat>}
625      %Equal.sym(List<&2, Nat>, visited(rows, extents, remaining, target, base, 1n+digit), shift(base, high),
626        visited_base(rows, extents, remaining, target, base, 1n+digit)) :
627        {List.append(&2, Nat, shift(Nat.add(base, offset), low), _) ==
628          shift(base, List.append(&2, Nat, visited(rest, others, span, inner, offset, first), high)) : List<&2, Nat>}
629      %Equal.sym(List<&2, Nat>, visited(rest, others, span, inner, offset, first), shift(offset, low),
630        visited_base(rest, others, span, inner, offset, first)) :
631        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) ==
632          shift(base, List.append(&2, Nat, _, high)) : List<&2, Nat>}
633      %Equal.sym(List<&2, Nat>, shift(base, List.append(&2, Nat, shift(offset, low), high)),
634        List.append(&2, Nat, shift(base, shift(offset, low)), shift(base, high)), shift_append(base, shift(offset, low), high)) :
635        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) == _ : List<&2, Nat>}
636      %Equal.sym(List<&2, Nat>, shift(base, shift(offset, low)), shift(Nat.add(base, offset), low), shift_shift(base, offset, low)) :
637        {List.append(&2, Nat, shift(Nat.add(base, offset), low), shift(base, high)) == List.append(&2, Nat, _, shift(base, high)) : List<&2, Nat>}
638      {==}
639    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
640      {==}