def plan_walk

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

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

67# The same for the rows of a plan's index map.
68def plan_walk(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
69  +exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>},
70  +target: Nat, +below: {Nat.is_lt(target, Shape.numel(Shape.extents(Shape.left_shape(plan)))) == True{} : Bool}) ->
71  Fiber.RestWalk(rows, Shape.extents(Shape.output_shape(plan)), target, Shape.numel(Shape.extents(Shape.output_shape(plan)))):
72  +left = Shape.extents(Shape.left_shape(plan))
73  +right = Shape.extents(Shape.right_shape(plan))
74  +output = Shape.extents(Shape.output_shape(plan))
75  %Equal.sym(List<&2, Shape.Row>, rows, Shape.index_table(left, right, output), exact) :
76    Fiber.RestWalk(_, output, target, Shape.numel(output))
77  table_walk(left, right, output, Shape.alignment(plan), target, below)