def cells_left

bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_PROOF.bend:61

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/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
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/tensor/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/reduce_to_fold_PROOF.bend as ReduceTo
import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws

61# `cells` reads the left input only at left_map(plan, k), which is in bounds,
62# so a left history that lists the same words gives the same cells.
63def cells_left(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
64  +left: Shape.Shape, +left_fill: F32, +kept: List<&2, StorageLaws.Op>, +left_ops: List<&2, StorageLaws.Op>,
65  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
66  +left_shape: {Shape.left_shape(plan) == left : Shape.Shape},
67  +same: {Contracts.words(left, left_fill, kept, Shape.logical_numel(Shape.dimensions(left))) ==
68    Contracts.words(left, left_fill, left_ops, Shape.logical_numel(Shape.dimensions(left))) : List<&2, F32>},
69  +count: Nat, +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
70  {Laws.cells(operation, plan, left, left_fill, kept, right, right_fill, right_ops, count) ==
71    Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, count) : List<&2, F32>}:
72  match count:
73    case 0n:
74      {==}
75    case 1n+previous:
76      +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
77      +below = Conversion.le_successor_lt(previous, total, room)
78      +index = U32.from_nat(PointwiseBroadcast.left_map(plan, previous))
79      +right_word = Replayed.value(right, right_fill, right_ops, U32.from_nat(PointwiseBroadcast.right_map(plan, previous)))
80      cells_snoc(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous,
81        Laws.cells(operation, plan, left, left_fill, kept, right, right_fill, right_ops, previous),
82        cell(operation, plan, left, left_fill, kept, right, right_fill, right_ops, previous),
83        cell(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous),
84        cells_left(operation, plan, left, left_fill, kept, left_ops, right, right_fill, right_ops, left_shape, same, previous,
85          Conversion.le_predecessor(previous, total, room)),
86        Equal.cong(F32, F32, word => Elementwise.scalar(operation, word, right_word),
87          Replayed.value(left, left_fill, kept, index), Replayed.value(left, left_fill, left_ops, index),
88          SumToSizeFoldProof.agree_at(left, left_fill, kept, left_ops, same, index,
89            PointwiseProof.word_inside(PointwiseBroadcast.left_map(plan, previous), left,
90              ReduceTo.target_below(plan, left, left_shape, previous, below)))))