def pair_step_lift

bendtorch-core@0.1.0.0 · src/tensor/contracts/fold_proof.bend:216

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold

216def pair_step_lift(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
217  +rest: Nat, +operation: O, +plan: P, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
218  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +cursor: Nat,
219  left_inside: {Nat.is_lt(U32.to_nat(U32.from_nat(ia(plan, cursor))), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool},
220  right_inside: {Nat.is_lt(U32.to_nat(U32.from_nat(ib(plan, cursor))), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool},
221  -Q: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32> -> Type,
222  next: PairStart(~O, ~g, ~P, ~ia, ~ib, rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, 1n+cursor, Q)) ->
223  PairStart(~O, ~g, ~P, ~ia, ~ib, 1n+rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, cursor, Q):
224  +left_seen = List.append(&2, StorageLaws.Op, left_ops, Pointwise.reads(~P, ~ia, plan, cursor))
225  +right_seen = List.append(&2, StorageLaws.Op, right_ops, Pointwise.reads(~P, ~ib, plan, cursor))
226  +left_index = U32.from_nat(ia(plan, cursor))
227  +right_index = U32.from_nat(ib(plan, cursor))
228  +before = Fold.pair_total(~O, ~g, ~P, ~ia, ~ib, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, cursor)
229  +x = Replayed.value(left, left_fill, left_seen, left_index)
230  left_fits => right_fits =>
231    Replayed.read_lift(left, left_fill, left_seen, left_fits, left_index, left_inside,
232      pair => Q(Fold.pair_left(~O, ~g, operation, right_index,
233        left => right => value => Fold.pair_loop(~O, ~g, ~P, ~ia, ~ib, rest, operation, plan, 1n+cursor, left, right, value),
234        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_seen), right_fits}}, before,
235        Core.read_dynamic_result(left, pair))),
236      left_after => Replayed.read_lift(right, right_fill, right_seen, right_fits, right_index, right_inside,
237        pair => Q(Fold.pair_right(~O, ~g, operation,
238          left => right => value => Fold.pair_loop(~O, ~g, ~P, ~ia, ~ib, rest, operation, plan, 1n+cursor, left, right, value),
239          Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_seen, StorageLaws.Read{left_index})), left_after}},
240          before, x, Core.read_dynamic_result(right, pair))),
241        right_after => pair_shift(~O, ~g, ~P, ~ia, ~ib, rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, cursor,
242          Q, next)(left_after, right_after)))