src/autograd/tensor_vjp_PROOF.bend

bendtorch-core@0.1.0.0 · 1,209 lines · 30 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
6import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
7import bendtorch-core@0.1.0.0/src/tensor/reduction.bend as Reduction
8import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
9import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
10import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
11import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
12import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
13import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
14import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_proofs.bend as BinaryProof
15# Proves binary_replayed, which importing the proof above would otherwise leave open.
16import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_PROOF.bend as BinaryLaw
17import bendtorch-core@0.1.0.0/src/autograd/reduce_to_fold_PROOF.bend as ReduceTo
18import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_LAWS.bend as SumToSizeFold
19import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
20import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
21import bendtorch-core@0.1.0.0/src/autograd/loss.bend as Loss
22import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws
23import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
24import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
25import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
26import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
27import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
28import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
29import bendtorch-core@0.1.0.0/src/tensor/permute_values.bend as PermuteValues
30import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as PermuteRecords
31import bendtorch-core@0.1.0.0/src/tensor/permute_values_contract_LAWS.bend as PermuteLaws
32import bendtorch-core@0.1.0.0/src/tensor/permute_values_contract_PROOF.bend as PermuteProof
33import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_proof.bend as MatmulProof
34import bendtorch-core@0.1.0.0/src/tensor/contracts/product_proof.bend as ProductProof
35import bendtorch-core@0.1.0.0/src/tensor/transpose_PROOF.bend as TransposeProof
36import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex
37
38def cell(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
39  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
40  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +index: Nat) -> F32:
41  Pointwise.cell(~Elementwise.Operation, ~Elementwise.scalar, ~Shape.BroadcastPlan, ~PointwiseBroadcast.left_map,
42    ~PointwiseBroadcast.right_map, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, index)
43
44def cells_snoc(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
45  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
46  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +previous: Nat,
47  +before: List<&2, F32>, +word: F32, +other: F32,
48  rest: {before == Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous) : List<&2, F32>},
49  last: {word == other : F32}) ->
50  {List.append(&2, F32, before, [word]) ==
51    List.append(&2, F32, Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous), [other]) : List<&2, F32>}:
52  Equal.trans(List<&2, F32>, List.append(&2, F32, before, [word]),
53    List.append(&2, F32, Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous), [word]),
54    List.append(&2, F32, Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous), [other]),
55    Equal.cong(List<&2, F32>, List<&2, F32>, values => List.append(&2, F32, values, [word]),
56      before, Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous), rest),
57    Equal.cong(F32, List<&2, F32>,
58      value => List.append(&2, F32, Laws.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous), [value]),
59      word, other, last))
60
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)))))
91
92def product_replayed(+operation: Elementwise.Operation, +left: Shape.Shape, +right: Shape.Shape, +plan: Shape.BroadcastPlan,
93  +planned: {Shape.broadcast_checked(left, right) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
94  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
95  -P: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type,
96  h: Laws.ProductWith(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, P)) ->
97  @-left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat} ->
98  @-right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat} ->
99    P(Elementwise.binary(operation, Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
100      Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})):
101  left_fits => right_fits =>
102    BinaryProof.binary_replayed(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, planned, P, h)(left_fits, right_fits)
103
104def Laws.tensor_product_replayed(operation, left, right, plan, planned, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits, P, h):
105  product_replayed(operation, left, right, plan, planned, left_fill, left_ops, right_fill, right_ops, P, h)(left_fits, right_fits)
106
107# Two delivered contributions, each a replayed tensor with its target's
108# dimensions whose cells hold the stated words, meet the pair contract.
109def pair_contract(+a: Nat, +left: Shape.Shape, -left_cell: U32 -> Maybe<&2, F32>,
110  +b: Nat, +right: Shape.Shape, -right_cell: U32 -> Maybe<&2, F32>, -saved: Backward.Saved, -g: Core.Dynamic,
111  +first: Shape.Shape, +first_dimensions: {Shape.dimensions(first) == Shape.dimensions(left) : List<&2, U32>},
112  +first_fill: F32, +first_ops: List<&2, StorageLaws.Op>,
113  first_values: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
114    {Some{Replayed.value(first, first_fill, first_ops, index)} == left_cell(index) : Maybe<&2, F32>},
115  -first_fits: {Lifecycle.logical_length(Replayed.storage(first, first_fill, first_ops)) == Shape.logical_numel(Shape.dimensions(first)) : Nat},
116  +second: Shape.Shape, +second_dimensions: {Shape.dimensions(second) == Shape.dimensions(right) : List<&2, U32>},
117  +second_fill: F32, +second_ops: List<&2, StorageLaws.Op>,
118  second_values: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
119    {Some{Replayed.value(second, second_fill, second_ops, index)} == right_cell(index) : Maybe<&2, F32>},
120  -second_fits: {Lifecycle.logical_length(Replayed.storage(second, second_fill, second_ops)) == Shape.logical_numel(Shape.dimensions(second)) : Nat}) ->
121  Laws.PairContract<a, left, left_cell, b, right, right_cell,
122    Backward.Pass{saved, g, Done{[Backward.Contribution{a, left, Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}},
123      Backward.Contribution{b, right, Core.Packed{second, Core.TTensor{Replayed.storage(second, second_fill, second_ops), second_fits}}}]}}>:
124  Laws.PairContract{
125    {==},
126    Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(first), Shape.dimensions(left),
127      first_dimensions),
128    index => inside =>
129      +inside = inside
130      Equal.trans(Maybe<&2, F32>,
131        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}, index)),
132        Some{Replayed.value(first, first_fill, first_ops, index)}, left_cell(index),
133        PointwiseProof.replayed_point(first, first_fill, first_ops, first_fits, index,
134          SumToSizeFoldProof.shape_below(first, left, index, first_dimensions, inside)),
135        first_values(index, inside)),
136    {==},
137    Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(second), Shape.dimensions(right),
138      second_dimensions),
139    index => inside =>
140      +inside = inside
141      Equal.trans(Maybe<&2, F32>,
142        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{second, Core.TTensor{Replayed.storage(second, second_fill, second_ops), second_fits}}, index)),
143        Some{Replayed.value(second, second_fill, second_ops, index)}, right_cell(index),
144        PointwiseProof.replayed_point(second, second_fill, second_ops, second_fits, index,
145          SumToSizeFoldProof.shape_below(second, right, index, second_dimensions, inside)),
146        second_values(index, inside))}
147
148def Laws.tensor_add_vjp_values(a, b, left, right, shape, left_plan, left_planned, left_same, right_plan, right_planned, right_same,
149  saved, fill, ops, fits):
150  +values = Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape)))
151  SumToSizeFold.sum_to_size_replayed(left, shape, left_plan, left_planned, left_same, fill, ops, fits, values, {==},
152    pair => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
153      b, right, index => Laws.reduced(right, shape, right_plan, values, index),
154      Backward.added_first(b, right, Backward.delivered_as(a, left, pair))>,
155    kept => kept_words => kept_fits => first => first_dimensions => first_fill => first_ops => first_values => first_fits =>
156      SumToSizeFold.sum_to_size_replayed(right, shape, right_plan, right_planned, right_same, fill, kept, kept_fits, values, kept_words,
157        pair => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
158          b, right, index => Laws.reduced(right, shape, right_plan, values, index),
159          Backward.added_second(Done{Backward.Contribution{a, left,
160            Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}}},
161            Backward.delivered_as(b, right, pair))>,
162        last => last_words => last_fits => second => second_dimensions => second_fill => second_ops => second_values => second_fits =>
163          pair_contract(a, left, index => Laws.reduced(left, shape, left_plan, values, index),
164            b, right, index => Laws.reduced(right, shape, right_plan, values, index),
165            Backward.Unsaved{}, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
166            first, first_dimensions, first_fill, first_ops, first_values, first_fits,
167            second, second_dimensions, second_fill, second_ops, second_values, second_fits)))
168
169def Laws.tensor_sub_vjp_values(a, b, left, right, shape, left_plan, left_planned, left_same, negated, negated_planned,
170  right_plan, right_planned, right_same, saved, fill, ops, fits):
171  +negated_planned = negated_planned
172  +count = Shape.logical_numel(Shape.dimensions(shape))
173  +values = Contracts.words(shape, fill, ops, count)
174  +zero = Reduction.scalar_shape()
175  +minus = Shape.output_shape(negated)
176  +minus_count = Shape.logical_numel(Shape.dimensions(minus))
177  +negatives = Laws.cells(Elementwise.NegateLeft{}, negated, shape, fill, ops, zero, 0.0, Nil{}, minus_count)
178  SumToSizeFold.sum_to_size_replayed(left, shape, left_plan, left_planned, left_same, fill, ops, fits, values, {==},
179    pair => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
180      b, right, index => Laws.reduced(right, minus, right_plan, negatives, index),
181      Backward.subtracted_first(b, right, Backward.delivered_as(a, left, pair))>,
182    kept => kept_words => kept_fits => first => first_dimensions => first_fill => first_ops => first_values => first_fits =>
183      -delivered = {Done{Backward.Contribution{a, left, Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}}} :
184        Result<&2, &1, Backward.Error, Backward.Contribution>}
185      Replayed.new_lift(zero, 0.0,
186        result => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
187          b, right, index => Laws.reduced(right, minus, right_plan, negatives, index),
188          Backward.subtracted_second(delivered, b, right,
189            Backward.negated(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}}, result))>,
190        zero_fits =>
191          product_replayed(Elementwise.NegateLeft{}, shape, zero, negated, negated_planned, fill, kept, 0.0, Nil{},
192            pair => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
193              b, right, index => Laws.reduced(right, minus, right_plan, negatives, index),
194              Backward.subtracted_second(delivered, b, right, Backward.negated_pair(pair))>,
195            last => last_words => last_fits => zero_kept => zero_words => zero_kept_fits => product_ops => product_values => product_fits =>
196              SumToSizeFold.sum_to_size_replayed(right, minus, right_plan, right_planned, right_same, 0.0, product_ops, product_fits, negatives,
197                Equal.trans(List<&2, F32>, Contracts.words(minus, 0.0, product_ops, minus_count),
198                  Laws.cells(Elementwise.NegateLeft{}, negated, shape, fill, kept, zero, 0.0, Nil{}, minus_count), negatives,
199                  product_values,
200                  cells_left(Elementwise.NegateLeft{}, negated, shape, fill, kept, ops, zero, 0.0, Nil{},
201                    ReduceTo.plan_left(shape, zero, negated, negated_planned), kept_words, minus_count, Conversion.le_reflexive(minus_count))),
202                pair => Laws.PairContract<a, left, index => Laws.reduced(left, shape, left_plan, values, index),
203                  b, right, index => Laws.reduced(right, minus, right_plan, negatives, index),
204                  Backward.Pass{Backward.Unsaved{}, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
205                    Backward.both(delivered, Backward.contribution_of(Backward.delivered_as(b, right, pair)))}>,
206                reduced => reduced_words => reduced_fits => second => second_dimensions => second_fill => second_ops => second_values => second_fits =>
207                  pair_contract(a, left, index => Laws.reduced(left, shape, left_plan, values, index),
208                    b, right, index => Laws.reduced(right, minus, right_plan, negatives, index),
209                    Backward.Unsaved{}, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
210                    first, first_dimensions, first_fill, first_ops, first_values, first_fits,
211                    second, second_dimensions, second_fill, second_ops, second_values, second_fits)))(kept_fits, zero_fits)))
212
213def Laws.tensor_mul_vjp_values(a, b, left, right, shape, first, first_planned, left_plan, left_planned, left_same,
214  second, second_planned, right_plan, right_planned, right_same, fill, ops, fits, left_fill, left_ops, left_fits,
215  right_fill, right_ops, right_fits):
216  +second_planned = second_planned
217  +front = Shape.output_shape(first)
218  +front_count = Shape.logical_numel(Shape.dimensions(front))
219  +back = Shape.output_shape(second)
220  +back_count = Shape.logical_numel(Shape.dimensions(back))
221  +front_values = Laws.cells(Elementwise.Mul{}, first, shape, fill, ops, right, right_fill, right_ops, front_count)
222  +back_values = Laws.cells(Elementwise.Mul{}, second, shape, fill, ops, left, left_fill, left_ops, back_count)
223  -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
224  product_replayed(Elementwise.Mul{}, shape, right, first, first_planned, fill, ops, right_fill, right_ops,
225    pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
226      b, right, index => Laws.reduced(right, back, right_plan, back_values, index),
227      Backward.multiplied_first(a, b, left, right, x, pair)>,
228    kept => kept_words => kept_fits => y => y_words => y_fits => front_ops => front_words => front_fits =>
229      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}} : Core.Dynamic}
230      -y_kept = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, y), y_fits}} : Core.Dynamic}
231      SumToSizeFold.sum_to_size_replayed(left, front, left_plan, left_planned, left_same, 0.0, front_ops, front_fits, front_values,
232        front_words,
233        pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
234          b, right, index => Laws.reduced(right, back, right_plan, back_values, index),
235          Backward.multiplied_second(b, right, Backward.contribution_of(Backward.delivered_as(a, left, pair)), y_kept,
236            Elementwise.mul(g, x))>,
237        reduced => reduced_words => reduced_fits => head => head_dimensions => head_fill => head_ops => head_values => head_fits =>
238          -delivered = {Done{Backward.Contribution{a, left, Core.Packed{head, Core.TTensor{Replayed.storage(head, head_fill, head_ops), head_fits}}}} :
239            Result<&2, &1, Backward.Error, Backward.Contribution>}
240          product_replayed(Elementwise.Mul{}, shape, left, second, second_planned, fill, kept, left_fill, left_ops,
241            pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
242              b, right, index => Laws.reduced(right, back, right_plan, back_values, index),
243              Backward.multiplied_second(b, right, delivered, y_kept, pair)>,
244            last => last_words => last_fits => x_kept => x_words => x_fits => back_ops => back_words => back_fits =>
245              SumToSizeFold.sum_to_size_replayed(right, back, right_plan, right_planned, right_same, 0.0, back_ops, back_fits, back_values,
246                Equal.trans(List<&2, F32>, Contracts.words(back, 0.0, back_ops, back_count),
247                  Laws.cells(Elementwise.Mul{}, second, shape, fill, kept, left, left_fill, left_ops, back_count), back_values,
248                  back_words,
249                  cells_left(Elementwise.Mul{}, second, shape, fill, kept, ops, left, left_fill, left_ops,
250                    ReduceTo.plan_left(shape, left, second, second_planned), kept_words, back_count, Conversion.le_reflexive(back_count))),
251                pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
252                  b, right, index => Laws.reduced(right, back, right_plan, back_values, index),
253                  Backward.Pass{Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, x_kept), x_fits}}, y_kept},
254                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
255                    Backward.both(delivered, Backward.contribution_of(Backward.delivered_as(b, right, pair)))}>,
256                tail => tail_words => tail_fits => second_out => second_dimensions => second_fill => second_ops => second_values => second_fits =>
257                  pair_contract(a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
258                    b, right, index => Laws.reduced(right, back, right_plan, back_values, index),
259                    Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, x_kept), x_fits}}, y_kept},
260                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
261                    head, head_dimensions, head_fill, head_ops, head_values, head_fits,
262                    second_out, second_dimensions, second_fill, second_ops, second_values, second_fits)))(kept_fits, left_fits)))(fits, right_fits)
263
264# A permuted copy holds, at every in-bounds index, the source's word at
265# permute_source: its words are `permuted_words`, which are the words of the
266# copy loop's history.
267def permuted_at(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
268  +output_ops: List<&2, StorageLaws.Op>,
269  +moved: {Contracts.words(Layout.output(input, axes, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) ==
270    PermuteRecords.permuted_words(input, fill, ops, axes, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) : List<&2, F32>},
271  +index: U32, +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
272  {Replayed.value(Layout.output(input, axes, plan), 0.0, output_ops, index) ==
273    Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(index)))) : F32}:
274  +output = Layout.output(input, axes, plan)
275  +total = Shape.logical_numel(Shape.dimensions(output))
276  +request : PermuteValues.Request = PermuteValues.Request{input, axes}
277  +history = IndexMap.writes(~PermuteValues.Request, ~PermuteValues.source_of, total, request, input, fill, ops)
278  Equal.trans(F32, Replayed.value(output, 0.0, output_ops, index), Replayed.value(output, 0.0, history, index),
279    Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(index)))),
280    SumToSizeFoldProof.agree_at(output, 0.0, output_ops, history,
281      Equal.trans(List<&2, F32>, Contracts.words(output, 0.0, output_ops, total), PermuteRecords.permuted_words(input, fill, ops, axes, total),
282        Contracts.words(output, 0.0, history, total), moved,
283        Equal.sym(List<&2, F32>, Contracts.words(output, 0.0, history, total), PermuteRecords.permuted_words(input, fill, ops, axes, total),
284          PermuteValues.words_moved(input, fill, ops, axes, plan, total, Conversion.le_reflexive(total)))),
285      index, inside),
286    PointwiseProof.writes_value(~Unit, ~IndexMap.keep, ~PermuteValues.Request, ~PermuteValues.source_of, ~PermuteValues.source_of, total, Unit{},
287      request, input, fill, ops, input, fill, ops, output, index, inside, Conversion.le_reflexive(total)))
288
289def factors_snoc(+before: List<&2, DotFold.Factors>, +after: List<&2, DotFold.Factors>, +x: F32, +y: F32, +u: F32, +v: F32,
290  rest: {before == after : List<&2, DotFold.Factors>}, xu: {x == u : F32}, yv: {y == v : F32}) ->
291  {List.append(&2, DotFold.Factors, before, [DotFold.Factors{x, y}]) ==
292    List.append(&2, DotFold.Factors, after, [DotFold.Factors{u, v}]) : List<&2, DotFold.Factors>}:
293  %rest : {List.append(&2, DotFold.Factors, before, [DotFold.Factors{x, y}]) ==
294    List.append(&2, DotFold.Factors, _, [DotFold.Factors{u, v}]) : List<&2, DotFold.Factors>}
295  %xu : {List.append(&2, DotFold.Factors, before, [DotFold.Factors{x, y}]) ==
296    List.append(&2, DotFold.Factors, before, [DotFold.Factors{_, v}]) : List<&2, DotFold.Factors>}
297  %yv : {List.append(&2, DotFold.Factors, before, [DotFold.Factors{x, y}]) ==
298    List.append(&2, DotFold.Factors, before, [DotFold.Factors{x, _}]) : List<&2, DotFold.Factors>}
299  {==}
300
301# The pairs the first product g @ y^T reads, on the copy of g and the
302# transposed copy of y, are the pairs right_transposed_factors names on g and
303# y themselves.
304def first_factors(+row: Nat, +column: Nat, +inner: U32, +right_inner: U32, +columns: U32,
305  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
306  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
307  +plan: Layout.Plan<right, Laws.swapped(right)>, +output_ops: List<&2, StorageLaws.Op>,
308  +moved: {Contracts.words(Laws.transposed(right, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Laws.transposed(right, plan)))) ==
309    PermuteRecords.permuted_words(right, right_fill, right_ops, Laws.swapped(right),
310      Shape.logical_numel(Shape.dimensions(Laws.transposed(right, plan)))) : List<&2, F32>},
311  +re: {Shape.dimensions(Laws.transposed(right, plan)) == [right_inner, columns] : List<&2, U32>},
312  +compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
313  +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
314  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
315  {Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, Replayed.snoc(ops, StorageLaws.Second{}),
316    Laws.transposed(right, plan), 0.0, output_ops, count) ==
317    Laws.right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
318      right, right_fill, right_ops, Laws.swapped(right), count) : List<&2, DotFold.Factors>}:
319  match count:
320    case 0n:
321      {==}
322    case 1n+previous:
323      +inner_count = U32.to_nat(inner)
324      +column_count = U32.to_nat(columns)
325      +axes = Laws.swapped(right)
326      +output = Laws.transposed(right, plan)
327      +position = Nat.add(Nat.mul(previous, column_count), column)
328      +below = MatmulProof.lt_bound(previous, inner_count, U32.to_nat(right_inner), compatible,
329        Conversion.le_successor_lt(previous, inner_count, room))
330      +bound = MatmulProof.cell_below(output, right_inner, columns, re, previous, below, column, column_below)
331      +cell = U32.from_nat(Nat.add(Nat.mul(row, inner_count), previous))
332      factors_snoc(
333        Contracts.factors(row, column, inner_count, column_count, shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), output, 0.0, output_ops, previous),
334        Laws.right_transposed_factors(row, column, inner_count, column_count, shape, fill, ops, right, right_fill, right_ops, axes, previous),
335        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), cell),
336        Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
337        Replayed.value(shape, fill, ops, cell),
338        Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, position))),
339        first_factors(row, column, inner, right_inner, columns, shape, fill, ops, right, right_fill, right_ops, plan, output_ops, moved,
340          re, compatible, column_below, previous, Conversion.le_predecessor(previous, inner_count, room)),
341        SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.Second{}, cell, {==}),
342        Equal.trans(F32, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
343          Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, U32.to_nat(U32.from_nat(position))))),
344          Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, position))),
345          permuted_at(right, right_fill, right_ops, axes, plan, output_ops, moved, U32.from_nat(position),
346            PointwiseProof.word_inside(position, output, bound)),
347          Equal.cong(Nat, F32, number => Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, number))),
348            U32.to_nat(U32.from_nat(position)), position, PointwiseProof.word_exact(position, output, bound))))
349
350# The pairs the second product x^T @ g reads, on the transposed copy of x and
351# the copy of g's first half, are the pairs left_transposed_factors names on x
352# and g themselves.
353def second_factors(+row: Nat, +column: Nat, +rows: U32, +inner: U32, +columns: U32,
354  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
355  +plan: Layout.Plan<left, Laws.swapped(left)>, +output_ops: List<&2, StorageLaws.Op>,
356  +moved: {Contracts.words(Laws.transposed(left, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Laws.transposed(left, plan)))) ==
357    PermuteRecords.permuted_words(left, left_fill, left_ops, Laws.swapped(left),
358      Shape.logical_numel(Shape.dimensions(Laws.transposed(left, plan)))) : List<&2, F32>},
359  +le: {Shape.dimensions(Laws.transposed(left, plan)) == [rows, inner] : List<&2, U32>},
360  +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
361  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
362  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
363  {Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), Laws.transposed(left, plan), 0.0, output_ops,
364    shape, fill, Replayed.snoc(Replayed.snoc(ops, StorageLaws.First{}), StorageLaws.Second{}), count) ==
365    Laws.left_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, Laws.swapped(left),
366      shape, fill, ops, count) : List<&2, DotFold.Factors>}:
367  match count:
368    case 0n:
369      {==}
370    case 1n+previous:
371      +inner_count = U32.to_nat(inner)
372      +column_count = U32.to_nat(columns)
373      +axes = Laws.swapped(left)
374      +output = Laws.transposed(left, plan)
375      +position = Nat.add(Nat.mul(row, inner_count), previous)
376      +bound = MatmulProof.cell_below(output, rows, inner, le, row, row_below, previous,
377        Conversion.le_successor_lt(previous, inner_count, room))
378      +cell = U32.from_nat(Nat.add(Nat.mul(previous, column_count), column))
379      +front = Replayed.snoc(ops, StorageLaws.First{})
380      factors_snoc(
381        Contracts.factors(row, column, inner_count, column_count, output, 0.0, output_ops, shape, fill, Replayed.snoc(front, StorageLaws.Second{}), previous),
382        Laws.left_transposed_factors(row, column, inner_count, column_count, left, left_fill, left_ops, axes, shape, fill, ops, previous),
383        Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
384        Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), cell),
385        Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, position))),
386        Replayed.value(shape, fill, ops, cell),
387        second_factors(row, column, rows, inner, columns, left, left_fill, left_ops, plan, output_ops, moved, le, row_below, shape, fill, ops,
388          previous, Conversion.le_predecessor(previous, inner_count, room)),
389        Equal.trans(F32, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
390          Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, U32.to_nat(U32.from_nat(position))))),
391          Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, position))),
392          permuted_at(left, left_fill, left_ops, axes, plan, output_ops, moved, U32.from_nat(position),
393            PointwiseProof.word_inside(position, output, bound)),
394          Equal.cong(Nat, F32, number => Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, number))),
395            U32.to_nat(U32.from_nat(position)), position, PointwiseProof.word_exact(position, output, bound))),
396        Equal.trans(F32, Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), cell), Replayed.value(shape, fill, front, cell),
397          Replayed.value(shape, fill, ops, cell),
398          SumToSizeFoldProof.half_keeps(shape, fill, front, StorageLaws.Second{}, cell, {==}),
399          SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.First{}, cell, {==})))
400
401# The record the Matmul VJP law concludes, for the plans the law names.
402def MatmulGoal(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
403  +right_plan: Layout.Plan<right, Laws.swapped(right)>, +first: Matrix.Plan<shape, Laws.transposed(right, right_plan)>,
404  +left_plan: Layout.Plan<left, Laws.swapped(left)>, +second: Matrix.Plan<Laws.transposed(left, left_plan), shape>,
405  +fill: F32, +ops: List<&2, StorageLaws.Op>, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
406  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, pass: Backward.Pass) -> Type:
407  +flipped = Laws.transposed(right, right_plan)
408  +turned = Laws.transposed(left, left_plan)
409  Laws.MatmulPairContract<a, left, Contracts.matrix_rows(shape, flipped, first), Contracts.matrix_columns(shape, flipped, first),
410    row => column => DotFold.dot_fold(Laws.right_transposed_factors(row, column,
411      Contracts.matrix_inner(shape, flipped, first), Contracts.matrix_columns(shape, flipped, first),
412      shape, fill, ops, right, right_fill, right_ops, Laws.swapped(right), Contracts.matrix_inner(shape, flipped, first)), 0.0),
413    b, right, Contracts.matrix_rows(turned, shape, second), Contracts.matrix_columns(turned, shape, second),
414    row => column => DotFold.dot_fold(Laws.left_transposed_factors(row, column,
415      Contracts.matrix_inner(turned, shape, second), Contracts.matrix_columns(turned, shape, second),
416      left, left_fill, left_ops, Laws.swapped(left), shape, fill, ops, Contracts.matrix_inner(turned, shape, second)), 0.0),
417    pass>
418
419# One product's cell: the contribution reads the replayed product, and the
420# product's cell is the ordered total over the copies, which first_factors or
421# second_factors turns into the stated pairs.
422def product_cell(+output: Shape.Shape, +rows: U32, +columns: U32, +oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
423  +output_ops: List<&2, StorageLaws.Op>,
424  -output_fits: {Lifecycle.logical_length(Replayed.storage(output, 0.0, output_ops)) == Shape.logical_numel(Shape.dimensions(output)) : Nat},
425  +row: Nat, +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
426  +column: Nat, +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
427  +read: List<&2, DotFold.Factors>, +stated: List<&2, DotFold.Factors>,
428  total: {Replayed.value(output, 0.0, output_ops, U32.from_nat(Nat.add(Nat.mul(row, U32.to_nat(columns)), column))) ==
429    DotFold.dot_fold(read, 0.0) : F32},
430  same: {read == stated : List<&2, DotFold.Factors>}) ->
431  {Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}},
432    U32.from_nat(Nat.add(Nat.mul(row, U32.to_nat(columns)), column)))) == Some{DotFold.dot_fold(stated, 0.0)} : Maybe<&2, F32>}:
433  +position = Nat.add(Nat.mul(row, U32.to_nat(columns)), column)
434  Equal.trans(Maybe<&2, F32>,
435    Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}},
436      U32.from_nat(position))),
437    Some{Replayed.value(output, 0.0, output_ops, U32.from_nat(position))}, Some{DotFold.dot_fold(stated, 0.0)},
438    PointwiseProof.replayed_point(output, 0.0, output_ops, output_fits, U32.from_nat(position),
439      PointwiseProof.word_inside(position, output, MatmulProof.cell_below(output, rows, columns, oe, row, row_below, column, column_below))),
440    Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
441      DotFold.dot_fold(stated, 0.0),
442      Equal.trans(F32, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)), DotFold.dot_fold(read, 0.0),
443        DotFold.dot_fold(stated, 0.0), total,
444        Equal.cong(List<&2, DotFold.Factors>, F32, list => DotFold.dot_fold(list, 0.0), read, stated, same))))
445
446def facts_left(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
447  facts: MatmulProof.MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}:
448  MatmulProof.MatrixFacts{le, re, oe, compatible} = facts
449  le
450
451def facts_right(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
452  facts: MatmulProof.MatrixFacts<left, right, output, rows, inner, right_inner, columns>) ->
453  {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}:
454  MatmulProof.MatrixFacts{le, re, oe, compatible} = facts
455  re
456
457def facts_compatible(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
458  facts: MatmulProof.MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}:
459  MatmulProof.MatrixFacts{le, re, oe, compatible} = facts
460  compatible
461
462def matmul_values(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
463  +right_inside: {True{} == Transpose.valid(right, 0, 1) : Bool},
464  +right_plan: Layout.Plan<right, Laws.swapped(right)>,
465  +right_laid: {Layout.permutation(right, Laws.swapped(right)) == Done{right_plan} : Result<&2, &2, Layout.Error, Layout.Plan<right, Laws.swapped(right)>>},
466  +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
467  -le: {Shape.dimensions(shape) == [rows, inner] : List<&2, U32>},
468  -re: {Shape.dimensions(Laws.transposed(right, right_plan)) == [right_inner, columns] : List<&2, U32>},
469  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
470  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
471  +first_planned: {Matrix.plan(shape, Laws.transposed(right, right_plan)) ==
472    Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
473    Result<&2, &2, Matrix.Error, Matrix.Plan<shape, Laws.transposed(right, right_plan)>>},
474  +first_same: {Shape.dimensions(output) == Shape.dimensions(left) : List<&2, U32>},
475  +left_inside: {True{} == Transpose.valid(left, 0, 1) : Bool},
476  +left_plan: Layout.Plan<left, Laws.swapped(left)>,
477  +left_laid: {Layout.permutation(left, Laws.swapped(left)) == Done{left_plan} : Result<&2, &2, Layout.Error, Layout.Plan<left, Laws.swapped(left)>>},
478  +back: Shape.Shape, +back_rows: U32, +back_inner: U32, +back_right_inner: U32, +back_columns: U32,
479  -back_le: {Shape.dimensions(Laws.transposed(left, left_plan)) == [back_rows, back_inner] : List<&2, U32>},
480  -back_re: {Shape.dimensions(shape) == [back_right_inner, back_columns] : List<&2, U32>},
481  -back_oe: {Shape.dimensions(back) == [back_rows, back_columns] : List<&2, U32>},
482  -back_compatible: {U32.to_nat(back_inner) == U32.to_nat(back_right_inner) : Nat},
483  +second_planned: {Matrix.plan(Laws.transposed(left, left_plan), shape) ==
484    Done{Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}} :
485    Result<&2, &2, Matrix.Error, Matrix.Plan<Laws.transposed(left, left_plan), shape>>},
486  +second_same: {Shape.dimensions(back) == Shape.dimensions(right) : List<&2, U32>},
487  +fill: F32, +ops: List<&2, StorageLaws.Op>,
488  -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
489  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
490  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
491  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
492  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
493  MatmulGoal(a, b, left, right, shape, right_plan, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible},
494    left_plan, Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
495    fill, ops, left_fill, left_ops, right_fill, right_ops,
496    Backward.vjp(Backward.Matmul{a, b, left, right},
497      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
498        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
499      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
500  +flipped = Laws.transposed(right, right_plan)
501  +turned = Laws.transposed(left, left_plan)
502  +first : Matrix.Plan<shape, Laws.transposed(right, right_plan)> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
503  +second : Matrix.Plan<Laws.transposed(left, left_plan), shape> =
504    Matrix.Plan{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}
505  +front_facts : MatmulProof.MatrixFacts<shape, flipped, output, rows, inner, right_inner, columns> =
506    MatmulProof.plan_facts(shape, flipped, first, first_planned)
507  +back_facts : MatmulProof.MatrixFacts<turned, shape, back, back_rows, back_inner, back_right_inner, back_columns> =
508    MatmulProof.plan_facts(turned, shape, second, second_planned)
509  +front_re = facts_right(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
510  +front_oe = MatmulProof.output_dimensions(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
511  +front_compatible = facts_compatible(shape, flipped, output, rows, inner, right_inner, columns, front_facts)
512  +turned_le = facts_left(turned, shape, back, back_rows, back_inner, back_right_inner, back_columns, back_facts)
513  +turned_oe = MatmulProof.output_dimensions(turned, shape, back, back_rows, back_inner, back_right_inner, back_columns, back_facts)
514  -Goal : Backward.Pass -> Type = pass => MatmulGoal(a, b, left, right, shape, right_plan, first, left_plan, second,
515    fill, ops, left_fill, left_ops, right_fill, right_ops, pass)
516  -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
517  -y = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}} : Core.Dynamic}
518  +front = Replayed.snoc(ops, StorageLaws.First{})
519  SumToSizeFoldProof.clone_lift(shape, fill, ops, fits,
520    pair => Goal(Backward.matmul_first(a, b, left, right, x, y, Core.clone_dynamic_result(shape, pair))),
521    g_fits => copy_fits =>
522      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, front), g_fits}} : Core.Dynamic}
523      -copy = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.Second{})), copy_fits}} : Core.Dynamic}
524      PermuteLaws.transpose_replayed(right, right_fill, right_ops, right_fits, 0, 1, right_inside, right_plan, right_laid,
525        pair => Goal(Backward.matmul_second(b, right, x, g, Backward.left_product(a, left, copy, pair))),
526        y_kept => y_words => y_fits => flipped_ops => flipped_moved => flipped_fits =>
527          -y_back = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, y_kept), y_fits}} : Core.Dynamic}
528          Contracts.matmul_replayed(shape, flipped, first, first_planned, fill, Replayed.snoc(ops, StorageLaws.Second{}), copy_fits,
529            0.0, flipped_ops, flipped_fits,
530            result => Goal(Backward.matmul_second(b, right, x, g, (y_back, Backward.multiplied(a, left, result)))),
531            copy_kept => copy_words => copy_kept_fits => flipped_kept => flipped_words => flipped_kept_fits => output_ops => front_cells => output_fits =>
532              -delivered = {Backward.Contribution{a, left, Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}} :
533                Backward.Contribution}
534              SumToSizeFoldProof.clone_lift(shape, fill, front, g_fits,
535                pair => Goal(Backward.matmul_third(b, right, Done{delivered}, y_back, x, Core.clone_dynamic_result(shape, pair))),
536                last_fits => again_fits =>
537                  -last = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.First{})), last_fits}} : Core.Dynamic}
538                  -again = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.Second{})), again_fits}} :
539                    Core.Dynamic}
540                  PermuteLaws.transpose_replayed(left, left_fill, left_ops, left_fits, 0, 1, left_inside, left_plan, left_laid,
541                    pair => Goal(Backward.matmul_fourth(Done{delivered}, y_back, last, Backward.right_product(b, right, again, pair))),
542                    x_kept => x_words => x_fits => turned_ops => turned_moved => turned_fits =>
543                      -x_back = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, x_kept), x_fits}} : Core.Dynamic}
544                      Contracts.matmul_replayed(turned, shape, second, second_planned, 0.0, turned_ops, turned_fits,
545                        fill, Replayed.snoc(front, StorageLaws.Second{}), again_fits,
546                        result => Goal(Backward.matmul_fourth(Done{delivered}, y_back, last, (x_back, Backward.multiplied(b, right, result)))),
547                        turned_kept => turned_words => turned_kept_fits => again_kept => again_words => again_kept_fits =>
548                          back_ops => back_cells => back_fits =>
549                          Laws.MatmulPairContract{
550                            {==},
551                            Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
552                              Shape.dimensions(output), Shape.dimensions(left), first_same),
553                            row => row_below => column => column_below =>
554                              +row_below = row_below
555                              +column_below = column_below
556                              product_cell(output, rows, columns, front_oe, output_ops, output_fits, row, row_below, column, column_below,
557                                Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill,
558                                  Replayed.snoc(ops, StorageLaws.Second{}), flipped, 0.0, flipped_ops, U32.to_nat(inner)),
559                                Laws.right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
560                                  right, right_fill, right_ops, Laws.swapped(right), U32.to_nat(inner)),
561                                front_cells(row, row_below, column, column_below),
562                                first_factors(row, column, inner, right_inner, columns, shape, fill, ops, right, right_fill, right_ops,
563                                  right_plan, flipped_ops, flipped_moved, front_re, front_compatible, column_below,
564                                  U32.to_nat(inner), Conversion.le_reflexive(U32.to_nat(inner)))),
565                            {==},
566                            Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
567                              Shape.dimensions(back), Shape.dimensions(right), second_same),
568                            row => row_below => column => column_below =>
569                              +row_below = row_below
570                              +column_below = column_below
571                              product_cell(back, back_rows, back_columns, turned_oe, back_ops, back_fits, row, row_below, column, column_below,
572                                Contracts.factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_columns), turned, 0.0, turned_ops,
573                                  shape, fill, Replayed.snoc(front, StorageLaws.Second{}), U32.to_nat(back_inner)),
574                                Laws.left_transposed_factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_columns),
575                                  left, left_fill, left_ops, Laws.swapped(left), shape, fill, ops, U32.to_nat(back_inner)),
576                                back_cells(row, row_below, column, column_below),
577                                second_factors(row, column, back_rows, back_inner, back_columns, left, left_fill, left_ops, left_plan,
578                                  turned_ops, turned_moved, turned_le, row_below, shape, fill, ops,
579                                  U32.to_nat(back_inner), Conversion.le_reflexive(U32.to_nat(back_inner))))}))))))
580
581def Laws.tensor_matmul_vjp_values(a, b, left, right, shape, right_inside, right_plan, right_laid, first, first_planned, first_same,
582  left_inside, left_plan, left_laid, second, second_planned, second_same, fill, ops, fits, left_fill, left_ops, left_fits,
583  right_fill, right_ops, right_fits):
584  match first:
585    case Matrix.Plan{+output, +rows, +inner, +right_inner, +columns, le, re, oe, compatible}:
586      match second:
587        case Matrix.Plan{+back, +back_rows, +back_inner, +back_right_inner, +back_columns, back_le, back_re, back_oe, back_compatible}:
588          matmul_values(a, b, left, right, shape, right_inside, right_plan, right_laid,
589            output, rows, inner, right_inner, columns, le, re, oe, compatible, first_planned, first_same,
590            left_inside, left_plan, left_laid,
591            back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible,
592            second_planned, second_same, fill, ops, fits, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits)
593
594# The pairs g @ W reads on the Second halves of the clones of g and W are
595# the pairs on g and W themselves.
596def nt_first_factors(+row: Nat, +column: Nat, +inner: Nat, +columns: Nat,
597  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
598  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +count: Nat) ->
599  {Contracts.factors(row, column, inner, columns, shape, fill, Replayed.snoc(ops, StorageLaws.Second{}),
600    right, right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}), count) ==
601    Contracts.factors(row, column, inner, columns, shape, fill, ops, right, right_fill, right_ops, count) : List<&2, DotFold.Factors>}:
602  match count:
603    case 0n:
604      {==}
605    case 1n+previous:
606      +left_cell = U32.from_nat(Nat.add(Nat.mul(row, inner), previous))
607      +right_cell = U32.from_nat(Nat.add(Nat.mul(previous, columns), column))
608      factors_snoc(
609        Contracts.factors(row, column, inner, columns, shape, fill, Replayed.snoc(ops, StorageLaws.Second{}),
610          right, right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}), previous),
611        Contracts.factors(row, column, inner, columns, shape, fill, ops, right, right_fill, right_ops, previous),
612        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), left_cell),
613        Replayed.value(right, right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}), right_cell),
614        Replayed.value(shape, fill, ops, left_cell),
615        Replayed.value(right, right_fill, right_ops, right_cell),
616        nt_first_factors(row, column, inner, columns, shape, fill, ops, right, right_fill, right_ops, previous),
617        SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.Second{}, left_cell, {==}),
618        SumToSizeFoldProof.half_keeps(right, right_fill, right_ops, StorageLaws.Second{}, right_cell, {==}))
619
620# The pairs g^T @ x reads on the Second half of the First half of g and the
621# Second half of x are the pairs on g and x themselves.
622def nt_second_factors(+row: Nat, +column: Nat, +inner: Nat, +rows: Nat, +columns: Nat,
623  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
624  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +count: Nat) ->
625  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
626    shape, fill, Replayed.snoc(Replayed.snoc(ops, StorageLaws.First{}), StorageLaws.Second{}),
627    left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), count) ==
628    Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
629      shape, fill, ops, left, left_fill, left_ops, count) : List<&2, DotFold.Factors>}:
630  match count:
631    case 0n:
632      {==}
633    case 1n+previous:
634      +front = Replayed.snoc(ops, StorageLaws.First{})
635      +g_cell = Matrix.read_index(Matrix.Swapped{}, inner, rows, row, previous)
636      +x_cell = Matrix.read_index(Matrix.Swapped{}, inner, columns, column, previous)
637      factors_snoc(
638        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
639          shape, fill, Replayed.snoc(front, StorageLaws.Second{}), left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), previous),
640        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
641          shape, fill, ops, left, left_fill, left_ops, previous),
642        Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), g_cell),
643        Replayed.value(left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), x_cell),
644        Replayed.value(shape, fill, ops, g_cell),
645        Replayed.value(left, left_fill, left_ops, x_cell),
646        nt_second_factors(row, column, inner, rows, columns, shape, fill, ops, left, left_fill, left_ops, previous),
647        Equal.trans(F32, Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), g_cell), Replayed.value(shape, fill, front, g_cell),
648          Replayed.value(shape, fill, ops, g_cell),
649          SumToSizeFoldProof.half_keeps(shape, fill, front, StorageLaws.Second{}, g_cell, {==}),
650          SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.First{}, g_cell, {==})),
651        SumToSizeFoldProof.half_keeps(left, left_fill, left_ops, StorageLaws.Second{}, x_cell, {==}))
652
653# The record the MatmulNT VJP law concludes, for the plans the law names.
654def NTGoal(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
655  +first: Matrix.Plan<shape, right>, +second: Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>,
656  +fill: F32, +ops: List<&2, StorageLaws.Op>, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
657  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, pass: Backward.Pass) -> Type:
658  Laws.MatmulPairContract<a, left, Contracts.matrix_rows(shape, right, first), Contracts.matrix_columns(shape, right, first),
659    row => column => DotFold.dot_fold(Contracts.factors(row, column,
660      Contracts.matrix_inner(shape, right, first), Contracts.matrix_columns(shape, right, first),
661      shape, fill, ops, right, right_fill, right_ops, Contracts.matrix_inner(shape, right, first)), 0.0),
662    b, right,
663    Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
664    Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
665    row => column => DotFold.dot_fold(Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{},
666      Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
667      Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
668      Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
669      shape, fill, ops, left, left_fill, left_ops, Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)), 0.0),
670    pass>
671
672def nt_values(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
673  +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
674  -le: {Shape.dimensions(shape) == [rows, inner] : List<&2, U32>},
675  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
676  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
677  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
678  +first_planned: {Matrix.plan(shape, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
679    Result<&2, &2, Matrix.Error, Matrix.Plan<shape, right>>},
680  +first_same: {Shape.dimensions(output) == Shape.dimensions(left) : List<&2, U32>},
681  +back: Shape.Shape, +back_rows: U32, +back_inner: U32, +back_right_inner: U32, +back_columns: U32,
682  -back_le: {Shape.dimensions(shape) == Matrix.arranged(Matrix.Swapped{}, back_rows, back_inner) : List<&2, U32>},
683  -back_re: {Shape.dimensions(left) == Matrix.arranged(Matrix.Swapped{}, back_columns, back_right_inner) : List<&2, U32>},
684  -back_oe: {Shape.dimensions(back) == [back_rows, back_columns] : List<&2, U32>},
685  -back_compatible: {U32.to_nat(back_inner) == U32.to_nat(back_right_inner) : Nat},
686  +second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) ==
687    Done{Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}} :
688    Result<&2, &2, Matrix.Error, Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>>},
689  +second_same: {Shape.dimensions(back) == Shape.dimensions(right) : List<&2, U32>},
690  +fill: F32, +ops: List<&2, StorageLaws.Op>,
691  -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
692  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
693  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
694  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
695  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
696  NTGoal(a, b, left, right, shape, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible},
697    Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
698    fill, ops, left_fill, left_ops, right_fill, right_ops,
699    Backward.vjp(Backward.MatmulNT{a, b, left, right},
700      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
701        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
702      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
703  +first : Matrix.Plan<shape, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
704  +second : Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left> =
705    Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}
706  +front_facts : MatmulProof.MatrixFacts<shape, right, output, rows, inner, right_inner, columns> =
707    MatmulProof.plan_facts(shape, right, first, first_planned)
708  +back_facts : ProductProof.ProductFacts<Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner, back_columns> =
709    ProductProof.plan_facts(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned)
710  +front_oe = MatmulProof.output_dimensions(shape, right, output, rows, inner, right_inner, columns, front_facts)
711  +back_oe_live = ProductProof.output_dimensions(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner,
712    back_columns, back_facts)
713  -Goal : Backward.Pass -> Type = pass => NTGoal(a, b, left, right, shape, first, second,
714    fill, ops, left_fill, left_ops, right_fill, right_ops, pass)
715  -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
716  -y = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}} : Core.Dynamic}
717  +front = Replayed.snoc(ops, StorageLaws.First{})
718  SumToSizeFoldProof.clone_lift(shape, fill, ops, fits,
719    pair => Goal(Backward.nt_first(a, b, left, right, x, y, Core.clone_dynamic_result(shape, pair))),
720    g_fits => copy_fits =>
721      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, front), g_fits}} : Core.Dynamic}
722      -copy = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.Second{})), copy_fits}} : Core.Dynamic}
723      SumToSizeFoldProof.clone_lift(right, right_fill, right_ops, right_fits,
724        pair => Goal(Backward.nt_second(b, right, x, g, Backward.nt_left(a, left, copy, Core.clone_dynamic_result(right, pair)))),
725        y_fits => flipped_fits =>
726          -y_back = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, Replayed.snoc(right_ops, StorageLaws.First{})), y_fits}} : Core.Dynamic}
727          Contracts.matmul_replayed(shape, right, first, first_planned, fill, Replayed.snoc(ops, StorageLaws.Second{}), copy_fits,
728            right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}), flipped_fits,
729            result => Goal(Backward.nt_second(b, right, x, g, (y_back, Backward.multiplied(a, left, result)))),
730            copy_kept => copy_words => copy_kept_fits => flipped_kept => flipped_words => flipped_kept_fits => output_ops => front_cells => output_fits =>
731              -delivered = {Backward.Contribution{a, left, Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}} :
732                Backward.Contribution}
733              SumToSizeFoldProof.clone_lift(shape, fill, front, g_fits,
734                pair => Goal(Backward.nt_third(b, right, Done{delivered}, y_back, x, Core.clone_dynamic_result(shape, pair))),
735                last_fits => again_fits =>
736                  -last = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.First{})), last_fits}} : Core.Dynamic}
737                  -again = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(front, StorageLaws.Second{})), again_fits}} :
738                    Core.Dynamic}
739                  SumToSizeFoldProof.clone_lift(left, left_fill, left_ops, left_fits,
740                    pair => Goal(Backward.nt_fourth(Done{delivered}, y_back, last, Backward.nt_right(b, right, again, Core.clone_dynamic_result(left, pair)))),
741                    x_fits => turned_fits =>
742                      -x_back = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_ops, StorageLaws.First{})), x_fits}} :
743                        Core.Dynamic}
744                      Contracts.product_replayed(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned,
745                        fill, Replayed.snoc(front, StorageLaws.Second{}), again_fits,
746                        left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), turned_fits,
747                        result => Goal(Backward.nt_fourth(Done{delivered}, y_back, last, (x_back, Backward.multiplied(b, right, result)))),
748                        again_kept => again_words => again_kept_fits => turned_kept => turned_words => turned_kept_fits =>
749                          back_ops => back_cells => back_fits =>
750                          Laws.MatmulPairContract{
751                            {==},
752                            Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
753                              Shape.dimensions(output), Shape.dimensions(left), first_same),
754                            row => row_below => column => column_below =>
755                              +row_below = row_below
756                              +column_below = column_below
757                              product_cell(output, rows, columns, front_oe, output_ops, output_fits, row, row_below, column, column_below,
758                                Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill,
759                                  Replayed.snoc(ops, StorageLaws.Second{}), right, right_fill, Replayed.snoc(right_ops, StorageLaws.Second{}),
760                                  U32.to_nat(inner)),
761                                Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
762                                  right, right_fill, right_ops, U32.to_nat(inner)),
763                                front_cells(row, row_below, column, column_below),
764                                nt_first_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
765                                  right, right_fill, right_ops, U32.to_nat(inner))),
766                            {==},
767                            Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
768                              Shape.dimensions(back), Shape.dimensions(right), second_same),
769                            row => row_below => column => column_below =>
770                              +row_below = row_below
771                              +column_below = column_below
772                              product_cell(back, back_rows, back_columns, back_oe_live, back_ops, back_fits, row, row_below, column, column_below,
773                                Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
774                                  U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, Replayed.snoc(front, StorageLaws.Second{}),
775                                  left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), U32.to_nat(back_inner)),
776                                Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
777                                  U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)),
778                                back_cells(row, row_below, column, column_below),
779                                nt_second_factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_rows), U32.to_nat(back_columns),
780                                  shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)))}))))))
781
782def Laws.tensor_matmul_nt_vjp_values(a, b, left, right, shape, first, first_planned, first_same, second, second_planned, second_same,
783  fill, ops, fits, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits):
784  match first:
785    case Matrix.Plan{+output, +rows, +inner, +right_inner, +columns, le, re, oe, compatible}:
786      match second:
787        case Matrix.Product{+back, +back_rows, +back_inner, +back_right_inner, +back_columns, back_le, back_re, back_oe, back_compatible}:
788          nt_values(a, b, left, right, shape, output, rows, inner, right_inner, columns, le, re, oe, compatible, first_planned, first_same,
789            back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible,
790            second_planned, second_same, fill, ops, fits, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits)
791
792# The pairs g^T @ x reads on the Second half of g and the Second half of x
793# are the pairs on g and x themselves.
794def frozen_factors(+row: Nat, +column: Nat, +inner: Nat, +rows: Nat, +columns: Nat,
795  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
796  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +count: Nat) ->
797  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
798    shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), count) ==
799    Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
800      shape, fill, ops, left, left_fill, left_ops, count) : List<&2, DotFold.Factors>}:
801  match count:
802    case 0n:
803      {==}
804    case 1n+previous:
805      +g_cell = Matrix.read_index(Matrix.Swapped{}, inner, rows, row, previous)
806      +x_cell = Matrix.read_index(Matrix.Swapped{}, inner, columns, column, previous)
807      factors_snoc(
808        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
809          shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), previous),
810        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
811          shape, fill, ops, left, left_fill, left_ops, previous),
812        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), g_cell),
813        Replayed.value(left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), x_cell),
814        Replayed.value(shape, fill, ops, g_cell),
815        Replayed.value(left, left_fill, left_ops, x_cell),
816        frozen_factors(row, column, inner, rows, columns, shape, fill, ops, left, left_fill, left_ops, previous),
817        SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.Second{}, g_cell, {==}),
818        SumToSizeFoldProof.half_keeps(left, left_fill, left_ops, StorageLaws.Second{}, x_cell, {==}))
819
820def FrozenGoal(+b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
821  +second: Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>,
822  +fill: F32, +ops: List<&2, StorageLaws.Op>, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, pass: Backward.Pass) -> Type:
823  Laws.MatmulSoleContract<b, right,
824    Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
825    Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
826    row => column => DotFold.dot_fold(Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{},
827      Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
828      Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
829      Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
830      shape, fill, ops, left, left_fill, left_ops, Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)), 0.0),
831    pass>
832
833def frozen_values(+b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape,
834  +back: Shape.Shape, +back_rows: U32, +back_inner: U32, +back_right_inner: U32, +back_columns: U32,
835  -back_le: {Shape.dimensions(shape) == Matrix.arranged(Matrix.Swapped{}, back_rows, back_inner) : List<&2, U32>},
836  -back_re: {Shape.dimensions(left) == Matrix.arranged(Matrix.Swapped{}, back_columns, back_right_inner) : List<&2, U32>},
837  -back_oe: {Shape.dimensions(back) == [back_rows, back_columns] : List<&2, U32>},
838  -back_compatible: {U32.to_nat(back_inner) == U32.to_nat(back_right_inner) : Nat},
839  +second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) ==
840    Done{Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}} :
841    Result<&2, &2, Matrix.Error, Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>>},
842  +second_same: {Shape.dimensions(back) == Shape.dimensions(right) : List<&2, U32>},
843  +fill: F32, +ops: List<&2, StorageLaws.Op>,
844  -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
845  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
846  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
847  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
848  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
849  FrozenGoal(b, left, right, shape,
850    Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
851    fill, ops, left_fill, left_ops,
852    Backward.vjp_matmul_nt_right(b, right,
853      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
854        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
855      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
856  +second : Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left> =
857    Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible}
858  +back_facts : ProductProof.ProductFacts<Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner, back_columns> =
859    ProductProof.plan_facts(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned)
860  +back_oe_live = ProductProof.output_dimensions(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, back, back_rows, back_inner, back_right_inner,
861    back_columns, back_facts)
862  -Goal : Backward.Pass -> Type = pass => FrozenGoal(b, left, right, shape, second, fill, ops, left_fill, left_ops, pass)
863  -x = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}} : Core.Dynamic}
864  -y = {Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}} : Core.Dynamic}
865  SumToSizeFoldProof.clone_lift(shape, fill, ops, fits,
866    pair => Goal(Backward.nt_alone_first(b, right, x, y, Core.clone_dynamic_result(shape, pair))),
867    g_fits => copy_fits =>
868      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.First{})), g_fits}} : Core.Dynamic}
869      -copy = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, Replayed.snoc(ops, StorageLaws.Second{})), copy_fits}} : Core.Dynamic}
870      SumToSizeFoldProof.clone_lift(left, left_fill, left_ops, left_fits,
871        pair => Goal(Backward.nt_alone_second(y, g, Backward.nt_right(b, right, copy, Core.clone_dynamic_result(left, pair)))),
872        x_fits => turned_fits =>
873          -x_back = {Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, Replayed.snoc(left_ops, StorageLaws.First{})), x_fits}} :
874            Core.Dynamic}
875          Contracts.product_replayed(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second, second_planned,
876            fill, Replayed.snoc(ops, StorageLaws.Second{}), copy_fits,
877            left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), turned_fits,
878            result => Goal(Backward.nt_alone_second(y, g, (x_back, Backward.multiplied(b, right, result)))),
879            again_kept => again_words => again_kept_fits => turned_kept => turned_words => turned_kept_fits =>
880              back_ops => back_cells => back_fits =>
881              Laws.MatmulSoleContract{
882                {==},
883                Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
884                  Shape.dimensions(back), Shape.dimensions(right), second_same),
885                row => row_below => column => column_below =>
886                  +row_below = row_below
887                  +column_below = column_below
888                  product_cell(back, back_rows, back_columns, back_oe_live, back_ops, back_fits, row, row_below, column, column_below,
889                    Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
890                      U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, Replayed.snoc(ops, StorageLaws.Second{}),
891                      left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), U32.to_nat(back_inner)),
892                    Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(back_inner),
893                      U32.to_nat(back_rows), U32.to_nat(back_columns), shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)),
894                    back_cells(row, row_below, column, column_below),
895                    frozen_factors(row, column, U32.to_nat(back_inner), U32.to_nat(back_rows), U32.to_nat(back_columns),
896                      shape, fill, ops, left, left_fill, left_ops, U32.to_nat(back_inner)))})))
897
898def same_refl(+number: Nat) -> {True{} == Nat.is_eq(number, number) : Bool}:
899  match number:
900    case 0n:
901      {==}
902    case 1n+rest:
903      same_refl(rest)
904
905# With a not b, Frozen{a, MatmulNT{a, b, ..}} runs the right product alone.
906def frozen_skips(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, -saved: Backward.Saved, -g: Core.Dynamic,
907  apart: {Nat.is_eq(a, b) == False{} : Bool}) ->
908  {Backward.vjp_matmul_nt_right(b, right, saved, g) == Backward.vjp(Backward.Frozen{a, Backward.MatmulNT{a, b, left, right}}, saved, g) :
909    Backward.Pass}:
910  %same_refl(a) : {Backward.vjp_matmul_nt_right(b, right, saved, g) ==
911    Backward.nt_frozen(a, a, b, left, right, saved, g, Bool.and(_, Bool.not(Nat.is_eq(a, b)))) : Backward.Pass}
912  %Equal.sym(Bool, Nat.is_eq(a, b), False{}, apart) : {Backward.vjp_matmul_nt_right(b, right, saved, g) ==
913    Backward.nt_frozen(a, a, b, left, right, saved, g, Bool.and(True{}, Bool.not(_))) : Backward.Pass}
914  {==}
915
916def Laws.tensor_matmul_nt_frozen_vjp_values(a, b, apart, left, right, shape, second, second_planned, second_same,
917  fill, ops, fits, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits):
918  match second:
919    case Matrix.Product{+back, +back_rows, +back_inner, +back_right_inner, +back_columns, back_le, back_re, back_oe, back_compatible}:
920      %frozen_skips(a, b, left, right,
921        Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
922          Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
923        Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, apart) :
924        FrozenGoal(b, left, right, shape,
925          Matrix.Product{back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible},
926          fill, ops, left_fill, left_ops, _)
927      frozen_values(b, left, right, shape, back, back_rows, back_inner, back_right_inner, back_columns, back_le, back_re, back_oe, back_compatible,
928        second_planned, second_same, fill, ops, fits, left_fill, left_ops, left_fits, right_fill, right_ops, right_fits)
929
930# The axes transpose_dynamic(t, 0, 1) permutes a rank-2 t by are [1, 0].
931def swapped_matrix(+shape: Shape.Shape, +a: U32, +b: U32, +matrix: {Shape.dimensions(shape) == [a, b] : List<&2, U32>}) ->
932  {Laws.swapped(shape) == [1, 0] : List<&2, U32>}:
933  %Equal.sym(List<&2, U32>, Shape.dimensions(shape), [a, b], matrix) :
934    {Transpose.requested_axes(Shape.rank_dimensions(_), U32.to_nat(0), U32.to_nat(1)) == [1, 0] : List<&2, U32>}
935  {==}
936
937# Factor k of cell (row, column) of a `Kept` right operand of dimensions
938# [columns, inner] is the word right^T holds at k * columns + column.
939def nt_word(+right: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +columns: U32, +inner: U32,
940  +matrix: {Shape.dimensions(right) == [columns, inner] : List<&2, U32>},
941  +column: Nat, +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
942  +k: Nat, +below: {Nat.is_lt(k, U32.to_nat(inner)) == True{} : Bool}) ->
943  {Replayed.value(right, fill, ops, U32.from_nat(Nat.add(Nat.mul(column, U32.to_nat(inner)), k))) ==
944    Replayed.value(right, fill, ops, U32.from_nat(Values.permute_source_index(right, Laws.swapped(right),
945      Nat.add(Nat.mul(k, U32.to_nat(columns)), column)))) : F32}:
946  +across = U32.to_nat(columns)
947  +span = U32.to_nat(inner)
948  +source = Nat.add(Nat.mul(column, span), k)
949  +extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(right), [columns, inner], matrix)
950  %Equal.sym(List<&2, U32>, Laws.swapped(right), [1, 0], swapped_matrix(right, columns, inner, matrix)) :
951    {Replayed.value(right, fill, ops, U32.from_nat(source)) ==
952      Replayed.value(right, fill, ops, U32.from_nat(Values.permute_source_index(right, _, Nat.add(Nat.mul(k, across), column)))) : F32}
953  %Equal.sym(Maybe<&2, U32>, Values.permute_source(right, [1, 0], U32.from_nat(Nat.add(Nat.mul(k, across), column))), Some{U32.from_nat(source)},
954    TransposeProof.permute_source_transposed(right, columns, inner, matrix, k, below, column, column_below)) :
955    {Replayed.value(right, fill, ops, U32.from_nat(source)) ==
956      Replayed.value(right, fill, ops, U32.from_nat(U32.to_nat(Maybe.default(&2, U32, _, 0)))) : F32}
957  %Equal.sym(Nat, U32.to_nat(U32.from_nat(source)), source,
958    MatrixIndex.word_exact(source, right, MatrixIndex.left_below(across, span, column, column_below, right, extents, k, below))) :
959    {Replayed.value(right, fill, ops, U32.from_nat(source)) == Replayed.value(right, fill, ops, U32.from_nat(_)) : F32}
960  {==}
961
962# Factor k of cell (row, column) of a `Swapped` left operand of dimensions
963# [inner, rows] is the word left^T holds at row * inner + k.
964def tn_word(+left: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +inner: U32, +rows: U32,
965  +matrix: {Shape.dimensions(left) == [inner, rows] : List<&2, U32>},
966  +row: Nat, +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
967  +k: Nat, +below: {Nat.is_lt(k, U32.to_nat(inner)) == True{} : Bool}) ->
968  {Replayed.value(left, fill, ops, U32.from_nat(Nat.add(Nat.mul(k, U32.to_nat(rows)), row))) ==
969    Replayed.value(left, fill, ops, U32.from_nat(Values.permute_source_index(left, Laws.swapped(left),
970      Nat.add(Nat.mul(row, U32.to_nat(inner)), k)))) : F32}:
971  +height = U32.to_nat(rows)
972  +span = U32.to_nat(inner)
973  +source = Nat.add(Nat.mul(k, height), row)
974  +extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(left), [inner, rows], matrix)
975  %Equal.sym(List<&2, U32>, Laws.swapped(left), [1, 0], swapped_matrix(left, inner, rows, matrix)) :
976    {Replayed.value(left, fill, ops, U32.from_nat(source)) ==
977      Replayed.value(left, fill, ops, U32.from_nat(Values.permute_source_index(left, _, Nat.add(Nat.mul(row, span), k)))) : F32}
978  %Equal.sym(Maybe<&2, U32>, Values.permute_source(left, [1, 0], U32.from_nat(Nat.add(Nat.mul(row, span), k))), Some{U32.from_nat(source)},
979    TransposeProof.permute_source_transposed(left, inner, rows, matrix, row, row_below, k, below)) :
980    {Replayed.value(left, fill, ops, U32.from_nat(source)) ==
981      Replayed.value(left, fill, ops, U32.from_nat(U32.to_nat(Maybe.default(&2, U32, _, 0)))) : F32}
982  %Equal.sym(Nat, U32.to_nat(U32.from_nat(source)), source,
983    MatrixIndex.word_exact(source, left, MatrixIndex.right_below(span, height, row, row_below, left, extents, k, below))) :
984    {Replayed.value(left, fill, ops, U32.from_nat(source)) == Replayed.value(left, fill, ops, U32.from_nat(_)) : F32}
985  {==}
986
987def nt_factors(+row: Nat, +column: Nat, +rows: Nat, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
988  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +columns: U32, +inner: U32,
989  +matrix: {Shape.dimensions(right) == [columns, inner] : List<&2, U32>},
990  +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
991  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
992  {Contracts.product_factors(row, column, Matrix.Kept{}, Matrix.Kept{}, U32.to_nat(inner), rows, U32.to_nat(columns),
993    left, left_fill, left_ops, right, right_fill, right_ops, count) ==
994    Laws.right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops,
995      right, right_fill, right_ops, Laws.swapped(right), count) : List<&2, DotFold.Factors>}:
996  match count:
997    case 0n:
998      {==}
999    case 1n+previous:
1000      +span = U32.to_nat(inner)
1001      +across = U32.to_nat(columns)
1002      factors_snoc(
1003        Contracts.product_factors(row, column, Matrix.Kept{}, Matrix.Kept{}, span, rows, across,
1004          left, left_fill, left_ops, right, right_fill, right_ops, previous),
1005        Laws.right_transposed_factors(row, column, span, across, left, left_fill, left_ops,
1006          right, right_fill, right_ops, Laws.swapped(right), previous),
1007        Replayed.value(left, left_fill, left_ops, U32.from_nat(Nat.add(Nat.mul(row, span), previous))),
1008        Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(column, span), previous))),
1009        Replayed.value(left, left_fill, left_ops, U32.from_nat(Nat.add(Nat.mul(row, span), previous))),
1010        Replayed.value(right, right_fill, right_ops,
1011          U32.from_nat(Values.permute_source_index(right, Laws.swapped(right), Nat.add(Nat.mul(previous, across), column)))),
1012        nt_factors(row, column, rows, left, left_fill, left_ops, right, right_fill, right_ops, columns, inner, matrix,
1013          column_below, previous, Conversion.le_predecessor(previous, span, room)),
1014        {==},
1015        nt_word(right, right_fill, right_ops, columns, inner, matrix, column, column_below, previous,
1016          Conversion.le_successor_lt(previous, span, room)))
1017
1018def tn_factors(+row: Nat, +column: Nat, +columns: Nat, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
1019  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +inner: U32, +rows: U32,
1020  +matrix: {Shape.dimensions(left) == [inner, rows] : List<&2, U32>},
1021  +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
1022  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
1023  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(inner), U32.to_nat(rows), columns,
1024    left, left_fill, left_ops, right, right_fill, right_ops, count) ==
1025    Laws.left_transposed_factors(row, column, U32.to_nat(inner), columns, left, left_fill, left_ops, Laws.swapped(left),
1026      right, right_fill, right_ops, count) : List<&2, DotFold.Factors>}:
1027  match count:
1028    case 0n:
1029      {==}
1030    case 1n+previous:
1031      +span = U32.to_nat(inner)
1032      +height = U32.to_nat(rows)
1033      factors_snoc(
1034        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, span, height, columns,
1035          left, left_fill, left_ops, right, right_fill, right_ops, previous),
1036        Laws.left_transposed_factors(row, column, span, columns, left, left_fill, left_ops, Laws.swapped(left),
1037          right, right_fill, right_ops, previous),
1038        Replayed.value(left, left_fill, left_ops, U32.from_nat(Nat.add(Nat.mul(previous, height), row))),
1039        Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(previous, columns), column))),
1040        Replayed.value(left, left_fill, left_ops,
1041          U32.from_nat(Values.permute_source_index(left, Laws.swapped(left), Nat.add(Nat.mul(row, span), previous)))),
1042        Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(previous, columns), column))),
1043        tn_factors(row, column, columns, left, left_fill, left_ops, right, right_fill, right_ops, inner, rows, matrix,
1044          row_below, previous, Conversion.le_predecessor(previous, span, room)),
1045        tn_word(left, left_fill, left_ops, inner, rows, matrix, row, row_below, previous,
1046          Conversion.le_successor_lt(previous, span, room)),
1047        {==})
1048
1049def Laws.nt_reads_transpose(row, column, rows, left, left_fill, left_ops, right, right_fill, right_ops, columns, inner, matrix,
1050  column_below, count, room):
1051  nt_factors(row, column, rows, left, left_fill, left_ops, right, right_fill, right_ops, columns, inner, matrix,
1052    column_below, count, room)
1053
1054def Laws.tn_reads_transpose(row, column, columns, left, left_fill, left_ops, right, right_fill, right_ops, inner, rows, matrix,
1055  row_below, count, room):
1056  tn_factors(row, column, columns, left, left_fill, left_ops, right, right_fill, right_ops, inner, rows, matrix,
1057    row_below, count, room)
1058
1059# A delivered contribution, a replayed tensor with its target's dimensions
1060# whose cells hold the stated words, meets the one-contribution contract.
1061def one_contract(+a: Nat, +target: Shape.Shape, -target_cell: U32 -> Maybe<&2, F32>, -saved: Backward.Saved, -g: Core.Dynamic,
1062  +first: Shape.Shape, +first_dimensions: {Shape.dimensions(first) == Shape.dimensions(target) : List<&2, U32>},
1063  +first_fill: F32, +first_ops: List<&2, StorageLaws.Op>,
1064  first_values: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
1065    {Some{Replayed.value(first, first_fill, first_ops, index)} == target_cell(index) : Maybe<&2, F32>},
1066  -first_fits: {Lifecycle.logical_length(Replayed.storage(first, first_fill, first_ops)) == Shape.logical_numel(Shape.dimensions(first)) : Nat}) ->
1067  Laws.OneContract<a, target, target_cell,
1068    Backward.Pass{saved, g, Done{[Backward.Contribution{a, target, Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}}]}}>:
1069  Laws.OneContract{
1070    {==},
1071    Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(first), Shape.dimensions(target),
1072      first_dimensions),
1073    index => inside =>
1074      +inside = inside
1075      Equal.trans(Maybe<&2, F32>,
1076        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}, index)),
1077        Some{Replayed.value(first, first_fill, first_ops, index)}, target_cell(index),
1078        PointwiseProof.replayed_point(first, first_fill, first_ops, first_fits, index,
1079          SumToSizeFoldProof.shape_below(first, target, index, first_dimensions, inside)),
1080        first_values(index, inside))}
1081
1082def Laws.tensor_relu_vjp_values(a, input, shape, plan, planned, input_plan, input_planned, input_same, zero, fill, ops, fits,
1083  input_fill, input_ops, input_fits):
1084  +front = Shape.output_shape(plan)
1085  +front_count = Shape.logical_numel(Shape.dimensions(front))
1086  +front_values = Laws.cells(Elementwise.ReluBackward{}, plan, input, input_fill, input_ops, shape, fill, ops, front_count)
1087  product_replayed(Elementwise.ReluBackward{}, input, shape, plan, planned, input_fill, input_ops, fill, ops,
1088    pair => Laws.OneContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1089      Backward.relu_passed(a, input, zero, pair)>,
1090    x_kept => x_words => x_fits => g_kept => g_words => g_fits => front_ops => front_words => front_fits =>
1091      -x = {Core.Packed{input, Core.TTensor{Replayed.storage(input, input_fill, x_kept), x_fits}} : Core.Dynamic}
1092      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, g_kept), g_fits}} : Core.Dynamic}
1093      SumToSizeFold.sum_to_size_replayed(input, front, input_plan, input_planned, input_same, 0.0, front_ops, front_fits, front_values,
1094        front_words,
1095        pair => Laws.OneContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1096          Backward.Pass{Backward.Operands{x, zero}, g, Backward.alone(Backward.contribution_of(Backward.delivered_as(a, input, pair)))}>,
1097        reduced => reduced_words => reduced_fits => head => head_dimensions => head_fill => head_ops => head_values => head_fits =>
1098          one_contract(a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1099            Backward.Operands{x, zero}, g, head, head_dimensions, head_fill, head_ops, head_values, head_fits)))(input_fits, fits)
1100
1101# The MseLoss VJP runs the Mul VJP's steps with MseGrad{2 / N} for Mul, d
1102# for the right operand and e for the left one, and both products with the
1103# same plan.
1104def Laws.tensor_mse_vjp_values(a, b, left, right, diff, shape, plan, planned, left_plan, left_planned, left_same,
1105  right_plan, right_planned, right_same, fill, ops, fits, e_fill, e_ops, e_fits, d_fill, d_ops, d_fits):
1106  +planned = planned
1107  +scale = Loss.mse_scale(diff)
1108  +operation = {Elementwise.MseGrad{scale} : Elementwise.Operation}
1109  +front = Shape.output_shape(plan)
1110  +front_count = Shape.logical_numel(Shape.dimensions(front))
1111  +front_values = Laws.cells(operation, plan, shape, fill, ops, diff, d_fill, d_ops, front_count)
1112  +back_values = Laws.cells(operation, plan, shape, fill, ops, diff, e_fill, e_ops, front_count)
1113  -x = {Core.Packed{diff, Core.TTensor{Replayed.storage(diff, e_fill, e_ops), e_fits}} : Core.Dynamic}
1114  product_replayed(operation, shape, diff, plan, planned, fill, ops, d_fill, d_ops,
1115    pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
1116      b, right, index => Laws.reduced(right, front, right_plan, back_values, index),
1117      Backward.mse_first(a, b, left, right, scale, x, pair)>,
1118    kept => kept_words => kept_fits => y => y_words => y_fits => front_ops => front_words => front_fits =>
1119      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}} : Core.Dynamic}
1120      -y_kept = {Core.Packed{diff, Core.TTensor{Replayed.storage(diff, d_fill, y), y_fits}} : Core.Dynamic}
1121      SumToSizeFold.sum_to_size_replayed(left, front, left_plan, left_planned, left_same, 0.0, front_ops, front_fits, front_values,
1122        front_words,
1123        pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
1124          b, right, index => Laws.reduced(right, front, right_plan, back_values, index),
1125          Backward.mse_second(b, right, Backward.contribution_of(Backward.delivered_as(a, left, pair)), y_kept,
1126            Elementwise.binary(operation, g, x))>,
1127        reduced => reduced_words => reduced_fits => head => head_dimensions => head_fill => head_ops => head_values => head_fits =>
1128          -delivered = {Done{Backward.Contribution{a, left, Core.Packed{head, Core.TTensor{Replayed.storage(head, head_fill, head_ops), head_fits}}}} :
1129            Result<&2, &1, Backward.Error, Backward.Contribution>}
1130          product_replayed(operation, shape, diff, plan, planned, fill, kept, e_fill, e_ops,
1131            pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
1132              b, right, index => Laws.reduced(right, front, right_plan, back_values, index),
1133              Backward.mse_second(b, right, delivered, y_kept, pair)>,
1134            last => last_words => last_fits => x_kept => x_words => x_fits => back_ops => back_words => back_fits =>
1135              SumToSizeFold.sum_to_size_replayed(right, front, right_plan, right_planned, right_same, 0.0, back_ops, back_fits, back_values,
1136                Equal.trans(List<&2, F32>, Contracts.words(front, 0.0, back_ops, front_count),
1137                  Laws.cells(operation, plan, shape, fill, kept, diff, e_fill, e_ops, front_count), back_values,
1138                  back_words,
1139                  cells_left(operation, plan, shape, fill, kept, ops, diff, e_fill, e_ops,
1140                    ReduceTo.plan_left(shape, diff, plan, planned), kept_words, front_count, Conversion.le_reflexive(front_count))),
1141                pair => Laws.PairContract<a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
1142                  b, right, index => Laws.reduced(right, front, right_plan, back_values, index),
1143                  Backward.Pass{Backward.Operands{Core.Packed{diff, Core.TTensor{Replayed.storage(diff, e_fill, x_kept), x_fits}}, y_kept},
1144                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
1145                    Backward.both(delivered, Backward.contribution_of(Backward.delivered_as(b, right, pair)))}>,
1146                tail => tail_words => tail_fits => second_out => second_dimensions => second_fill => second_ops => second_values => second_fits =>
1147                  pair_contract(a, left, index => Laws.reduced(left, front, left_plan, front_values, index),
1148                    b, right, index => Laws.reduced(right, front, right_plan, back_values, index),
1149                    Backward.Operands{Core.Packed{diff, Core.TTensor{Replayed.storage(diff, e_fill, x_kept), x_fits}}, y_kept},
1150                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
1151                    head, head_dimensions, head_fill, head_ops, head_values, head_fits,
1152                    second_out, second_dimensions, second_fill, second_ops, second_values, second_fits)))(kept_fits, e_fits)))(fits, d_fits)
1153
1154# The CrossEntropy VJP runs the Mul VJP's steps with OneHotGrad for the
1155# first product (of g and the one-hot rows) and SoftmaxGrad for the second
1156# (of g and lp), both delivered to a with the input's shape.
1157def Laws.tensor_cross_entropy_vjp_values(a, input, shape, plan, planned, input_plan, input_planned, input_same, fill, ops, fits,
1158  lp_fill, lp_ops, lp_fits, hot_fill, hot_ops, hot_fits):
1159  +planned = planned
1160  +input_planned = input_planned
1161  +input_same = input_same
1162  +count = Loss.count_f32(Loss.rows(input))
1163  +first_operation = {Elementwise.OneHotGrad{count} : Elementwise.Operation}
1164  +single = Loss.one_class(input)
1165  +second_operation = {Elementwise.SoftmaxGrad{count, single} : Elementwise.Operation}
1166  +front = Shape.output_shape(plan)
1167  +front_count = Shape.logical_numel(Shape.dimensions(front))
1168  +front_values = Laws.cells(first_operation, plan, shape, fill, ops, input, hot_fill, hot_ops, front_count)
1169  +back_values = Laws.cells(second_operation, plan, shape, fill, ops, input, lp_fill, lp_ops, front_count)
1170  -x = {Core.Packed{input, Core.TTensor{Replayed.storage(input, lp_fill, lp_ops), lp_fits}} : Core.Dynamic}
1171  product_replayed(first_operation, shape, input, plan, planned, fill, ops, hot_fill, hot_ops,
1172    pair => Laws.PairContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1173      a, input, index => Laws.reduced(input, front, input_plan, back_values, index),
1174      Backward.entropy_first(a, input, count, single, x, pair)>,
1175    kept => kept_words => kept_fits => y => y_words => y_fits => front_ops => front_words => front_fits =>
1176      -g = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}} : Core.Dynamic}
1177      -y_kept = {Core.Packed{input, Core.TTensor{Replayed.storage(input, hot_fill, y), y_fits}} : Core.Dynamic}
1178      SumToSizeFold.sum_to_size_replayed(input, front, input_plan, input_planned, input_same, 0.0, front_ops, front_fits, front_values,
1179        front_words,
1180        pair => Laws.PairContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1181          a, input, index => Laws.reduced(input, front, input_plan, back_values, index),
1182          Backward.entropy_second(a, input, Backward.contribution_of(Backward.delivered_as(a, input, pair)), y_kept,
1183            Elementwise.binary(second_operation, g, x))>,
1184        reduced => reduced_words => reduced_fits => head => head_dimensions => head_fill => head_ops => head_values => head_fits =>
1185          -delivered = {Done{Backward.Contribution{a, input, Core.Packed{head, Core.TTensor{Replayed.storage(head, head_fill, head_ops), head_fits}}}} :
1186            Result<&2, &1, Backward.Error, Backward.Contribution>}
1187          product_replayed(second_operation, shape, input, plan, planned, fill, kept, lp_fill, lp_ops,
1188            pair => Laws.PairContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1189              a, input, index => Laws.reduced(input, front, input_plan, back_values, index),
1190              Backward.entropy_second(a, input, delivered, y_kept, pair)>,
1191            last => last_words => last_fits => x_kept => x_words => x_fits => back_ops => back_words => back_fits =>
1192              SumToSizeFold.sum_to_size_replayed(input, front, input_plan, input_planned, input_same, 0.0, back_ops, back_fits, back_values,
1193                Equal.trans(List<&2, F32>, Contracts.words(front, 0.0, back_ops, front_count),
1194                  Laws.cells(second_operation, plan, shape, fill, kept, input, lp_fill, lp_ops, front_count), back_values,
1195                  back_words,
1196                  cells_left(second_operation, plan, shape, fill, kept, ops, input, lp_fill, lp_ops,
1197                    ReduceTo.plan_left(shape, input, plan, planned), kept_words, front_count, Conversion.le_reflexive(front_count))),
1198                pair => Laws.PairContract<a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1199                  a, input, index => Laws.reduced(input, front, input_plan, back_values, index),
1200                  Backward.Pass{Backward.Operands{Core.Packed{input, Core.TTensor{Replayed.storage(input, lp_fill, x_kept), x_fits}}, y_kept},
1201                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
1202                    Backward.both(delivered, Backward.contribution_of(Backward.delivered_as(a, input, pair)))}>,
1203                tail => tail_words => tail_fits => second_out => second_dimensions => second_fill => second_ops => second_values => second_fits =>
1204                  pair_contract(a, input, index => Laws.reduced(input, front, input_plan, front_values, index),
1205                    a, input, index => Laws.reduced(input, front, input_plan, back_values, index),
1206                    Backward.Operands{Core.Packed{input, Core.TTensor{Replayed.storage(input, lp_fill, x_kept), x_fits}}, y_kept},
1207                    Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, last), last_fits}},
1208                    head, head_dimensions, head_fill, head_ops, head_values, head_fits,
1209                    second_out, second_dimensions, second_fill, second_ops, second_values, second_fits)))(kept_fits, lp_fits)))(fits, hot_fits)