def table_walk

bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:50

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/autograd/reduce_fiber_proof.bend as Fiber

50# The index table of an aligned broadcast is walked as `Fiber.fiber_walk`
51# says.
52def table_walk(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>,
53  +target: Nat, +below: {Nat.is_lt(target, Shape.numel(left)) == True{} : Bool}) ->
54  Fiber.RestWalk(Shape.index_table(left, right, output), output, target, Shape.numel(output)):
55  match alignment:
56    case Shape.BAlign{+lp, +rp, minimal, +aligned}:
57      +nl = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, left))
58      +nr = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, right))
59      Fiber.fiber_walk(Shape.pad(nl, left), Shape.pad(nr, right), output,
60        %padding_left(left, right, output, lp, rp, aligned) : Shape.Aligned(Shape.pad(_, left), Shape.pad(nr, right), output)
61        %padding_right(left, right, output, lp, rp, aligned) : Shape.Aligned(Shape.pad(lp, left), Shape.pad(_, right), output)
62        aligned,
63        target,
64        %Equal.sym(Nat, Shape.numel(Shape.pad(nl, left)), Shape.numel(left), numel_pad(nl, left)) : {Nat.is_lt(target, _) == True{} : Bool}
65        below)