def cells_step

bendtorch-core@0.1.0.0 · src/autograd/broadcast_pointwise_PROOF.bend:50

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/storage/arith.bend as Arith
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_proof.bend as PointwiseProof
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_PROOF.bend as FoldProof
import bendtorch-core@0.1.0.0/src/autograd/broadcast_pointwise_LAWS.bend as Laws

50# Word `position` of `cells` is the input word at left_map(plan, position).
51def cells_step(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +position: Nat,
52  below: {Nat.is_lt(position, 1n+previous) == True{} : Bool}, decision: Shape.BoolCase<Nat.is_eq(position, previous)>,
53  earlier: @inside: {Nat.is_lt(position, previous) == True{} : Bool} ->
54    {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, previous), position) ==
55      Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}) ->
56  {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, 1n+previous), position) ==
57    Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}:
58  match decision:
59    case Shape.BFalse{equal}:
60      +before = Laws.cells(plan, shape, fill, ops, previous)
61      +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))
62      +inside = PointwiseProof.lt_succ_ne(position, previous, below, equal)
63      Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(List.append(&2, F32, before, [last]), position), ReduceLaws.word_at(before, position),
64        Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))},
65        FoldProof.word_at_front(before, last, position,
66          %Equal.sym(Nat, List.length(&2, F32, before), previous, cells_length(plan, shape, fill, ops, previous)) :
67            {Nat.is_lt(position, _) == True{} : Bool}
68          inside),
69        earlier(inside))
70    case Shape.BTrue{equal}:
71      +before = Laws.cells(plan, shape, fill, ops, previous)
72      +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))
73      %Equal.sym(Nat, position, previous, Arith.eq_true(position, previous, Equal.sym(Bool, True{}, Nat.is_eq(position, previous), equal))) :
74        {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) ==
75          Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, _)))} : Maybe<&2, F32>}
76      %cells_length(plan, shape, fill, ops, previous) :
77        {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{last} : Maybe<&2, F32>}
78      FoldProof.word_at_last(before, last)