def offsets_walk

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as BroadcastMap

79# As in reduce_fold_PROOF.bend: below the output count, the specified
80# offsets are the index walks.
81def offsets_walk(+plan: Shape.BroadcastPlan,
82  +map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
83  +position: Nat, +bound: {True{} == Nat.is_lt(position, Shape.elements(Shape.output_shape(plan))) : Bool}) ->
84  {Projection.offset_spec(plan, position) ==
85    Some{(U32.from_nat(BroadcastMap.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
86        Shape.extents(Shape.output_shape(plan)), map, position)),
87      U32.from_nat(BroadcastMap.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
88        Shape.extents(Shape.output_shape(plan)), map, position)))} : Maybe<&1, U32 & U32>}:
89  match map:
90    case Shape.Indexed{+rows, +exact}:
91      +left = Shape.extents(Shape.left_shape(plan))
92      +right = Shape.extents(Shape.right_shape(plan))
93      +output = Shape.extents(Shape.output_shape(plan))
94      +index = {Shape.FIndex{position, bound} : Shape.FlatIndex<Shape.output_shape(plan)>}
95      +slot = ShapeLaws.fin_of(Shape.numel(output), position)
96      +x = U32.from_nat(Shape.walk_left(rows, position))
97      +y = U32.from_nat(Shape.walk_right(rows, position))
98      %ShapeLaws.project_left_index_refines(plan, Shape.Indexed{rows, exact}, index) :
99        {Projection.pair_offsets(_, ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), slot))) ==
100          Some{(x, y)} : Maybe<&1, U32 & U32>}
101      %ShapeLaws.project_right_index_refines(plan, Shape.Indexed{rows, exact}, index) :
102        {Projection.pair_offsets(Some{x}, _) == Some{(x, y)} : Maybe<&1, U32 & U32>}
103      {==}
104    case Shape.Unindexed{zero}:
105      +output = Shape.output_shape(plan)
106      Empty.absurd({Projection.offset_spec(plan, position) == Some{(U32.from_nat(0n), U32.from_nat(0n))} : Maybe<&1, U32 & U32>},
107        Arith.lt_zero(position,
108          %Equal.trans(Nat, Shape.elements(output), Shape.numel(Shape.extents(output)), 0n, Shape.elements_exact(output),
109            Equal.sym(Nat, 0n, Shape.numel(Shape.extents(output)), zero)) : {Nat.is_lt(position, _) == True{} : Bool}
110          Equal.sym(Bool, True{}, Nat.is_lt(position, Shape.elements(output)), bound)))