def step_lift
bendtorch-core@0.1.0.0 · src/tensor/contracts/pointwise_proof.bend:337
import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
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
337def step_lift(~O: Data, ~f: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
338 +rest: Nat, +operation: O, +plan: P, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
339 +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +shape: Shape.Shape, +cursor: Nat,
340 left_inside: {Nat.is_lt(U32.to_nat(U32.from_nat(ia(plan, cursor))), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool},
341 right_inside: {Nat.is_lt(U32.to_nat(U32.from_nat(ib(plan, cursor))), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool},
342 output_inside: {Nat.is_lt(U32.to_nat(U32.from_nat(cursor)), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool},
343 -Q: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type,
344 next: Start(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, shape, 1n+cursor, Q)) ->
345 Start(~O, ~f, ~P, ~ia, ~ib, 1n+rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, shape, cursor, Q):
346 +left_seen = List.append(&2, StorageLaws.Op, left_ops, Pointwise.reads(~P, ~ia, plan, cursor))
347 +right_seen = List.append(&2, StorageLaws.Op, right_ops, Pointwise.reads(~P, ~ib, plan, cursor))
348 +history = Pointwise.writes(~O, ~f, ~P, ~ia, ~ib, cursor, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops)
349 +left_index = U32.from_nat(ia(plan, cursor))
350 +right_index = U32.from_nat(ib(plan, cursor))
351 +x = Replayed.value(left, left_fill, left_seen, left_index)
352 +y = Replayed.value(right, right_fill, right_seen, right_index)
353 left_fits => right_fits => output_fits =>
354 Replayed.read_lift(left, left_fill, left_seen, left_fits, left_index, left_inside,
355 pair => Q(Pointwise.left_step(~O, ~f, operation, shape, U32.from_nat(cursor), right_index,
356 state => Pointwise.loop(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, shape, 1n+cursor, state),
357 Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_seen), right_fits}},
358 Core.TTensor{Replayed.storage(shape, 0.0, history), output_fits}, Core.read_dynamic_result(left, pair))),
359 left_after => Replayed.read_lift(right, right_fill, right_seen, right_fits, right_index, right_inside,
360 pair => Q(Pointwise.right_step(~O, ~f, operation, shape, U32.from_nat(cursor),
361 state => Pointwise.loop(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, shape, 1n+cursor, state),
362 Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_seen, StorageLaws.Read{left_index})), left_after}},
363 Core.TTensor{Replayed.storage(shape, 0.0, history), output_fits}, x, Core.read_dynamic_result(right, pair))),
364 right_after => Replayed.write_lift(shape, 0.0, history, output_fits, U32.from_nat(cursor), f(operation, x, y), output_inside,
365 pair => Q(Pointwise.write_step(shape,
366 state => Pointwise.loop(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, shape, 1n+cursor, state),
367 Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_seen, StorageLaws.Read{left_index})), left_after}},
368 Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, Replayed.snoc(right_seen, StorageLaws.Read{right_index})), right_after}},
369 pair)),
370 output_after => shift(~O, ~f, ~P, ~ia, ~ib, rest, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, shape, cursor,
371 Q, next)(left_after, right_after, output_after))))