def step_lift

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

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

65def step_lift(~O: Data, ~g: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~keep: P -> Nat -> Bool,
66  +rest: Nat, +operation: O, +plan: P, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +cursor: Nat,
67  inside: {Nat.is_lt(U32.to_nat(U32.from_nat(ia(plan, cursor))), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool},
68  -Q: Core.Dynamic & Result<&2, &1, Core.Error, F32> -> Type,
69  next: Start(~O, ~g, ~P, ~ia, ~keep, rest, operation, plan, shape, fill, ops, 1n+cursor, Q)) ->
70  Start(~O, ~g, ~P, ~ia, ~keep, 1n+rest, operation, plan, shape, fill, ops, cursor, Q):
71  +seen = List.append(&2, StorageLaws.Op, ops, Pointwise.reads(~P, ~ia, plan, cursor))
72  +before = Fold.total(~O, ~g, ~P, ~ia, ~keep, operation, plan, shape, fill, ops, cursor)
73  fits =>
74    Replayed.read_lift(shape, fill, seen, fits, U32.from_nat(ia(plan, cursor)), inside,
75      pair => Q(Fold.step(~O, ~g, operation, keep(plan, cursor),
76        input => value => Fold.loop(~O, ~g, ~P, ~ia, ~keep, rest, operation, plan, 1n+cursor, input, value),
77        before, Core.read_dynamic_result(shape, pair))),
78      after => shift(~O, ~g, ~P, ~ia, ~keep, rest, operation, plan, shape, fill, ops, cursor, Q, next)(after))