def offsets_walk

bendtorch-core@0.1.0.0 · src/autograd/reduce_fold_PROOF.bend:26

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
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/tensor/contracts/broadcast.bend as Broadcast

26def offsets_walk(+plan: Shape.BroadcastPlan,
27  +map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
28  +position: Nat, +bound: {True{} == Nat.is_lt(position, Shape.elements(Shape.output_shape(plan))) : Bool}) ->
29  {Projection.offset_spec(plan, position) ==
30    Some{(U32.from_nat(Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
31        Shape.extents(Shape.output_shape(plan)), map, position)),
32      U32.from_nat(Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
33        Shape.extents(Shape.output_shape(plan)), map, position)))} : Maybe<&1, U32 & U32>}:
34  match map:
35    case Shape.Indexed{+rows, +exact}:
36      +left = Shape.extents(Shape.left_shape(plan))
37      +right = Shape.extents(Shape.right_shape(plan))
38      +output = Shape.extents(Shape.output_shape(plan))
39      +index = {Shape.FIndex{position, bound} : Shape.FlatIndex<Shape.output_shape(plan)>}
40      +slot = ShapeLaws.fin_of(Shape.numel(output), position)
41      +x = U32.from_nat(Shape.walk_left(rows, position))
42      +y = U32.from_nat(Shape.walk_right(rows, position))
43      %ShapeLaws.project_left_index_refines(plan, Shape.Indexed{rows, exact}, index) :
44        {Projection.pair_offsets(_, ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), slot))) ==
45          Some{(x, y)} : Maybe<&1, U32 & U32>}
46      %ShapeLaws.project_right_index_refines(plan, Shape.Indexed{rows, exact}, index) :
47        {Projection.pair_offsets(Some{x}, _) == Some{(x, y)} : Maybe<&1, U32 & U32>}
48      {==}
49    case Shape.Unindexed{zero}:
50      +output = Shape.output_shape(plan)
51      Empty.absurd({Projection.offset_spec(plan, position) == Some{(U32.from_nat(0n), U32.from_nat(0n))} : Maybe<&1, U32 & U32>},
52        Arith.lt_zero(position,
53          %Equal.trans(Nat, Shape.elements(output), Shape.numel(Shape.extents(output)), 0n, Shape.elements_exact(output),
54            Equal.sym(Nat, 0n, Shape.numel(Shape.extents(output)), zero)) : {Nat.is_lt(position, _) == True{} : Bool}
55          Equal.sym(Bool, True{}, Nat.is_lt(position, Shape.elements(output)), bound)))