src/autograd/tensor_vjp_LAWS.bend

bendtorch-core@0.1.0.0 · 626 lines · 25 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/layout.bend as Layout
6import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
7import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
8import bendtorch-core@0.1.0.0/src/tensor/reduction.bend as Reduction
9import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
10import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
11import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
12import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
13import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
14import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
15import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_LAWS.bend as Binary
16import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_LAWS.bend as SumToSizeFold
17import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
18import bendtorch-core@0.1.0.0/src/autograd/loss.bend as Loss
19import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as TensorBackward
20
21# The values of the tensor VJPs, stated through the pointwise, ordered-fold
22# and index-map copy families: each law takes a replayed gradient (and replayed saved operands),
23# and states, for both contributions, the target, the dimensions and the word
24# at every in-bounds position.
25
26# The words of a pointwise product in index order: word k is
27# scalar(operation, left[left_map(plan, k)], right[right_map(plan, k)]), the
28# value the pointwise family's broadcast instance writes at k
29# (`binary_LAWS.bend` in the tensor contracts).
30def cells(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
31  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
32  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +count: Nat) -> List<&2, F32>:
33  Binary.cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, count)
34
35# What sum_to_size(contribution, target) holds at target cell index, for a
36# contribution of `shape` whose words are `values`: the word itself when the
37# target has the contribution's dimensions, the ordered fiber sum otherwise.
38def reduced(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, +values: List<&2, F32>, +index: U32) ->
39  Maybe<&2, F32>:
40  SumToSizeFold.summed(Layout.words_same(Shape.dimensions(shape), Shape.dimensions(target)), plan,
41    Shape.logical_numel(Shape.dimensions(shape)), values, index)
42
43# What the transport form of tensor_product_replayed asks of P: P holds for
44# both inputs, replayed with any histories that keep their words, next to the
45# output replayed over +0.0 with any history whose words are `cells`.
46def ProductWith(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
47  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
48  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
49  -P: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type) -> Type:
50  Binary.BinaryWith(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, P)
51
52# The pointwise product the Sub and Mul VJPs compute, `Elementwise.binary`
53# with its parallel build, writes the words of the pointwise family's broadcast
54# instance, in transport form: for two replayed inputs whose shapes broadcast,
55# for every P, if P holds for every result that hands back both inputs with
56# their words and gives an output of the plan's shape whose words are `cells`,
57# P holds for the product. The output is a replayed tensor, so
58# sum_to_size_replayed applies to it. This is the tensor law binary_replayed
59# stated for the backward.
60law tensor_product_replayed:
61  for +operation: Elementwise.Operation
62  for +left: Shape.Shape
63  for +right: Shape.Shape
64  for +plan: Shape.BroadcastPlan
65  for planned: {Shape.broadcast_checked(left, right) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
66  for +left_fill: F32
67  for +left_ops: List<&2, StorageLaws.Op>
68  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
69  for +right_fill: F32
70  for +right_ops: List<&2, StorageLaws.Op>
71  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
72  for -P: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type
73  for h: ProductWith(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, P)
74  P(Elementwise.binary(operation, Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
75    Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}))
76
77# The two contributions of a VJP, when it sends exactly two.
78def first_of(pass: Backward.Pass) -> Maybe<&1, Backward.Contribution>:
79  Backward.Pass{saved, g, contributions} = pass
80  match contributions:
81    case Fail{error}:
82      None{}
83    case Done{Nil{}}:
84      None{}
85    case Done{Con{first, Nil{}}}:
86      None{}
87    case Done{Con{first, Con{second, Nil{}}}}:
88      Some{first}
89    case Done{Con{first, Con{second, Con{other, rest}}}}:
90      None{}
91
92def second_of(pass: Backward.Pass) -> Maybe<&1, Backward.Contribution>:
93  Backward.Pass{saved, g, contributions} = pass
94  match contributions:
95    case Fail{error}:
96      None{}
97    case Done{Nil{}}:
98      None{}
99    case Done{Con{first, Nil{}}}:
100      None{}
101    case Done{Con{first, Con{second, Nil{}}}}:
102      Some{second}
103    case Done{Con{first, Con{second, Con{other, rest}}}}:
104      None{}
105
106# What a VJP that sends two contributions promises: the first goes to node a
107# with left's dimensions and holds left_cell(i) at every in-bounds i, the
108# second goes to node b with right's dimensions and holds right_cell(i).
109type PairContract<-a: Nat, -left: Shape.Shape, -left_cell: U32 -> Maybe<&2, F32>,
110  -b: Nat, -right: Shape.Shape, -right_cell: U32 -> Maybe<&2, F32>, -pass: Backward.Pass> is Type:
111  PairContract{
112    first_target: {TensorBackward.sole_target(first_of(pass)) == Some{a} : Maybe<&2, Nat>},
113    first_dimensions: {TensorBackward.sole_dimensions(first_of(pass)) == Some{Shape.dimensions(left)} : Maybe<&2, List<&2, U32>>},
114    first_value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
115      {TensorBackward.sole_point(first_of(pass), index) == left_cell(index) : Maybe<&2, F32>},
116    second_target: {TensorBackward.sole_target(second_of(pass)) == Some{b} : Maybe<&2, Nat>},
117    second_dimensions: {TensorBackward.sole_dimensions(second_of(pass)) == Some{Shape.dimensions(right)} : Maybe<&2, List<&2, U32>>},
118    second_value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
119      {TensorBackward.sole_point(second_of(pass), index) == right_cell(index) : Maybe<&2, F32>}}
120
121# d(a+b)/da and d(a+b)/db: for a replayed gradient g, and operand shapes that
122# each broadcast to exactly g's shape, the Add VJP sends g reduced to the left
123# operand's shape to a, then g reduced to the right operand's shape to b. Each
124# cell holds g's own word when the operand has g's dimensions (-0.0 stays
125# -0.0), and otherwise the words of g that land on it, added in index order
126# from +0.0.
127law tensor_add_vjp_values:
128  for +a: Nat
129  for +b: Nat
130  for +left: Shape.Shape
131  for +right: Shape.Shape
132  for +shape: Shape.Shape
133  for +left_plan: Shape.BroadcastPlan
134  for left_planned: {Shape.broadcast_checked(left, shape) == Done{left_plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
135  for left_same: {Shape.dimensions(Shape.output_shape(left_plan)) == Shape.dimensions(shape) : List<&2, U32>}
136  for +right_plan: Shape.BroadcastPlan
137  for right_planned: {Shape.broadcast_checked(right, shape) == Done{right_plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
138  for right_same: {Shape.dimensions(Shape.output_shape(right_plan)) == Shape.dimensions(shape) : List<&2, U32>}
139  for saved: Backward.Saved
140  for +fill: F32
141  for +ops: List<&2, StorageLaws.Op>
142  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
143  PairContract<a, left, index => reduced(left, shape, left_plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index),
144    b, right, index => reduced(right, shape, right_plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index),
145    Backward.vjp(Backward.Add{a, b, left, right}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
146
147# d(a-b)/da and d(a-b)/db: the Sub VJP sends g reduced to the left operand's
148# shape to a, as Add does, then -g reduced to the right operand's shape to b.
149# -g is the pointwise product of g with NegateLeft against a rank-0 +0.0, so
150# its word k is F32.neg of g's word at left_map(negated, k): the sign bit is
151# flipped before any word is added, as in PyTorch.
152law tensor_sub_vjp_values:
153  for +a: Nat
154  for +b: Nat
155  for +left: Shape.Shape
156  for +right: Shape.Shape
157  for +shape: Shape.Shape
158  for +left_plan: Shape.BroadcastPlan
159  for left_planned: {Shape.broadcast_checked(left, shape) == Done{left_plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
160  for left_same: {Shape.dimensions(Shape.output_shape(left_plan)) == Shape.dimensions(shape) : List<&2, U32>}
161  for +negated: Shape.BroadcastPlan
162  for negated_planned: {Shape.broadcast_checked(shape, Reduction.scalar_shape()) == Done{negated} :
163    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
164  for +right_plan: Shape.BroadcastPlan
165  for right_planned: {Shape.broadcast_checked(right, Shape.output_shape(negated)) == Done{right_plan} :
166    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
167  for right_same: {Shape.dimensions(Shape.output_shape(right_plan)) == Shape.dimensions(Shape.output_shape(negated)) : List<&2, U32>}
168  for saved: Backward.Saved
169  for +fill: F32
170  for +ops: List<&2, StorageLaws.Op>
171  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
172  PairContract<a, left, index => reduced(left, shape, left_plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index),
173    b, right, index => reduced(right, Shape.output_shape(negated), right_plan,
174      cells(Elementwise.NegateLeft{}, negated, shape, fill, ops, Reduction.scalar_shape(), 0.0, Nil{},
175        Shape.logical_numel(Shape.dimensions(Shape.output_shape(negated)))), index),
176    Backward.vjp(Backward.Sub{a, b, left, right}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
177
178# d(a*b)/da and d(a*b)/db: with the saved operand values x (left) and y
179# (right), the Mul VJP sends g * y reduced to the left operand's shape to a,
180# then g * x reduced to the right operand's shape to b. Word k of g * y is
181# F32.mul(g[left_map(first, k)], y[right_map(first, k)]), and of g * x
182# F32.mul(g[left_map(second, k)], x[right_map(second, k)]).
183law tensor_mul_vjp_values:
184  for +a: Nat
185  for +b: Nat
186  for +left: Shape.Shape
187  for +right: Shape.Shape
188  for +shape: Shape.Shape
189  for +first: Shape.BroadcastPlan
190  for first_planned: {Shape.broadcast_checked(shape, right) == Done{first} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
191  for +left_plan: Shape.BroadcastPlan
192  for left_planned: {Shape.broadcast_checked(left, Shape.output_shape(first)) == Done{left_plan} :
193    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
194  for left_same: {Shape.dimensions(Shape.output_shape(left_plan)) == Shape.dimensions(Shape.output_shape(first)) : List<&2, U32>}
195  for +second: Shape.BroadcastPlan
196  for second_planned: {Shape.broadcast_checked(shape, left) == Done{second} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
197  for +right_plan: Shape.BroadcastPlan
198  for right_planned: {Shape.broadcast_checked(right, Shape.output_shape(second)) == Done{right_plan} :
199    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
200  for right_same: {Shape.dimensions(Shape.output_shape(right_plan)) == Shape.dimensions(Shape.output_shape(second)) : List<&2, U32>}
201  for +fill: F32
202  for +ops: List<&2, StorageLaws.Op>
203  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
204  for +left_fill: F32
205  for +left_ops: List<&2, StorageLaws.Op>
206  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
207  for +right_fill: F32
208  for +right_ops: List<&2, StorageLaws.Op>
209  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
210  PairContract<a, left, index => reduced(left, Shape.output_shape(first), left_plan,
211      cells(Elementwise.Mul{}, first, shape, fill, ops, right, right_fill, right_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(first)))),
212      index),
213    b, right, index => reduced(right, Shape.output_shape(second), right_plan,
214      cells(Elementwise.Mul{}, second, shape, fill, ops, left, left_fill, left_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(second)))),
215      index),
216    Backward.vjp(Backward.Mul{a, b, left, right},
217      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
218        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
219      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
220
221# The axes transpose_dynamic(t, 0, 1) permutes t by, and the shape of the
222# copy the plan gives.
223def swapped(+shape: Shape.Shape) -> List<&2, U32>:
224  Transpose.swapped_axes(shape, 0, 1)
225
226def transposed(+shape: Shape.Shape, +plan: Layout.Plan<shape, swapped(shape)>) -> Shape.Shape:
227  Layout.output(shape, swapped(shape), plan)
228
229# The first `count` factor pairs of cell (row, column) of left @ right^T,
230# where right^T is right permuted by `axes`: left[row * inner + k] and the
231# word of right at permute_source(right, axes, k * columns + column), in
232# increasing k. Contracts.factors reads the same pairs from the copy.
233def right_transposed_factors(+row: Nat, +column: Nat, +inner: Nat, +columns: Nat,
234  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
235  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +count: Nat) ->
236  List<&2, DotFold.Factors>:
237  match count:
238    case 0n:
239      Nil{}
240    case 1n+previous:
241      List.append(&2, DotFold.Factors,
242        right_transposed_factors(row, column, inner, columns, left, left_fill, left_ops, right, right_fill, right_ops, axes, previous),
243        [DotFold.Factors{Replayed.value(left, left_fill, left_ops, U32.from_nat(Nat.add(Nat.mul(row, inner), previous))),
244          Replayed.value(right, right_fill, right_ops,
245            U32.from_nat(Values.permute_source_index(right, axes, Nat.add(Nat.mul(previous, columns), column))))}])
246
247# The first `count` factor pairs of cell (row, column) of left^T @ right,
248# where left^T is left permuted by `axes`: the word of left at
249# permute_source(left, axes, row * inner + k) and right[k * columns + column],
250# in increasing k.
251def left_transposed_factors(+row: Nat, +column: Nat, +inner: Nat, +columns: Nat,
252  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>,
253  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +count: Nat) -> List<&2, DotFold.Factors>:
254  match count:
255    case 0n:
256      Nil{}
257    case 1n+previous:
258      List.append(&2, DotFold.Factors,
259        left_transposed_factors(row, column, inner, columns, left, left_fill, left_ops, axes, right, right_fill, right_ops, previous),
260        [DotFold.Factors{Replayed.value(left, left_fill, left_ops,
261            U32.from_nat(Values.permute_source_index(left, axes, Nat.add(Nat.mul(row, inner), previous)))),
262          Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(previous, columns), column)))}])
263
264# What the Matmul VJP promises: the first contribution goes to node a with
265# left's dimensions and holds left_cell(row, column) at row * left_columns +
266# column for every row < left_rows and column < left_columns; the second goes
267# to node b with right's dimensions and holds right_cell(row, column) likewise.
268type MatmulPairContract<-a: Nat, -left: Shape.Shape, -left_rows: Nat, -left_columns: Nat, -left_cell: Nat -> Nat -> F32,
269  -b: Nat, -right: Shape.Shape, -right_rows: Nat, -right_columns: Nat, -right_cell: Nat -> Nat -> F32, -pass: Backward.Pass> is Type:
270  MatmulPairContract{
271    first_target: {TensorBackward.sole_target(first_of(pass)) == Some{a} : Maybe<&2, Nat>},
272    first_dimensions: {TensorBackward.sole_dimensions(first_of(pass)) == Some{Shape.dimensions(left)} : Maybe<&2, List<&2, U32>>},
273    first_value: @+row: Nat -> @row_below: {Nat.is_lt(row, left_rows) == True{} : Bool} ->
274      @+column: Nat -> @column_below: {Nat.is_lt(column, left_columns) == True{} : Bool} ->
275      {TensorBackward.sole_point(first_of(pass), U32.from_nat(Nat.add(Nat.mul(row, left_columns), column))) ==
276        Some{left_cell(row, column)} : Maybe<&2, F32>},
277    second_target: {TensorBackward.sole_target(second_of(pass)) == Some{b} : Maybe<&2, Nat>},
278    second_dimensions: {TensorBackward.sole_dimensions(second_of(pass)) == Some{Shape.dimensions(right)} : Maybe<&2, List<&2, U32>>},
279    second_value: @+row: Nat -> @row_below: {Nat.is_lt(row, right_rows) == True{} : Bool} ->
280      @+column: Nat -> @column_below: {Nat.is_lt(column, right_columns) == True{} : Bool} ->
281      {TensorBackward.sole_point(second_of(pass), U32.from_nat(Nat.add(Nat.mul(row, right_columns), column))) ==
282        Some{right_cell(row, column)} : Maybe<&2, F32>}}
283
284# d(a@b)/da and d(a@b)/db: with the saved operand values x (left) and y
285# (right), the Matmul VJP sends g @ y^T to a, then x^T @ g to b. y^T and x^T
286# are transpose_dynamic(., 0, 1), the copies whose plans right_plan and
287# left_plan name; first and second are the plans of the two products. Cell
288# (row, column) of the first holds the ordered dot product of g's row with the
289# words of y that the transposed column reads, added in increasing k from
290# +0.0, and likewise for the second; each contribution has its operand's
291# dimensions. Nothing is reduced: the products already have them.
292law tensor_matmul_vjp_values:
293  for +a: Nat
294  for +b: Nat
295  for +left: Shape.Shape
296  for +right: Shape.Shape
297  for +shape: Shape.Shape
298  for right_inside: {True{} == Transpose.valid(right, 0, 1) : Bool}
299  for +right_plan: Layout.Plan<right, swapped(right)>
300  for right_laid: {Layout.permutation(right, swapped(right)) == Done{right_plan} : Result<&2, &2, Layout.Error, Layout.Plan<right, swapped(right)>>}
301  for +first: Matrix.Plan<shape, transposed(right, right_plan)>
302  for first_planned: {Matrix.plan(shape, transposed(right, right_plan)) == Done{first} :
303    Result<&2, &2, Matrix.Error, Matrix.Plan<shape, transposed(right, right_plan)>>}
304  for first_same: {Shape.dimensions(Matrix.output_shape(shape, transposed(right, right_plan), first)) == Shape.dimensions(left) : List<&2, U32>}
305  for left_inside: {True{} == Transpose.valid(left, 0, 1) : Bool}
306  for +left_plan: Layout.Plan<left, swapped(left)>
307  for left_laid: {Layout.permutation(left, swapped(left)) == Done{left_plan} : Result<&2, &2, Layout.Error, Layout.Plan<left, swapped(left)>>}
308  for +second: Matrix.Plan<transposed(left, left_plan), shape>
309  for second_planned: {Matrix.plan(transposed(left, left_plan), shape) == Done{second} :
310    Result<&2, &2, Matrix.Error, Matrix.Plan<transposed(left, left_plan), shape>>}
311  for second_same: {Shape.dimensions(Matrix.output_shape(transposed(left, left_plan), shape, second)) == Shape.dimensions(right) : List<&2, U32>}
312  for +fill: F32
313  for +ops: List<&2, StorageLaws.Op>
314  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
315  for +left_fill: F32
316  for +left_ops: List<&2, StorageLaws.Op>
317  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
318  for +right_fill: F32
319  for +right_ops: List<&2, StorageLaws.Op>
320  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
321  MatmulPairContract<a, left,
322    Contracts.matrix_rows(shape, transposed(right, right_plan), first), Contracts.matrix_columns(shape, transposed(right, right_plan), first),
323    row => column => DotFold.dot_fold(right_transposed_factors(row, column,
324      Contracts.matrix_inner(shape, transposed(right, right_plan), first), Contracts.matrix_columns(shape, transposed(right, right_plan), first),
325      shape, fill, ops, right, right_fill, right_ops, swapped(right), Contracts.matrix_inner(shape, transposed(right, right_plan), first)), 0.0),
326    b, right,
327    Contracts.matrix_rows(transposed(left, left_plan), shape, second), Contracts.matrix_columns(transposed(left, left_plan), shape, second),
328    row => column => DotFold.dot_fold(left_transposed_factors(row, column,
329      Contracts.matrix_inner(transposed(left, left_plan), shape, second), Contracts.matrix_columns(transposed(left, left_plan), shape, second),
330      left, left_fill, left_ops, swapped(left), shape, fill, ops, Contracts.matrix_inner(transposed(left, left_plan), shape, second)), 0.0),
331    Backward.vjp(Backward.Matmul{a, b, left, right},
332      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
333        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
334      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
335
336# d(x @ W^T)/dx and d(x @ W^T)/dW, the backward of Linear's product: with the
337# saved operand values x (left, [batch, in]) and W (right, [out, in], PyTorch's
338# layout) and the gradient g ([batch, out]), the MatmulNT VJP sends g @ W to
339# a, then g^T @ x to b. Neither product builds a transpose: first is the plan
340# of `matmul(g, W)`, and second the plan of `matmul_tn(g, x)`, which reads g
341# in place. Cell (i, l) of the first holds the ordered dot product over j of
342# g[i, j] * W[j, l], the words of g at i * out + j and of W at j * in + l;
343# cell (j, l) of the second holds the ordered dot product over i of
344# g[i, j] * x[i, l], the words of g at i * out + j and of x at i * in + l.
345# Each sum adds its terms in increasing order from +0.0, and each
346# contribution has its operand's dimensions. The bias of Linear is an Add
347# entry, whose VJP `tensor_add_vjp_values` states: its gradient is g summed
348# over the batch, column by column.
349law tensor_matmul_nt_vjp_values:
350  for +a: Nat
351  for +b: Nat
352  for +left: Shape.Shape
353  for +right: Shape.Shape
354  for +shape: Shape.Shape
355  for +first: Matrix.Plan<shape, right>
356  for first_planned: {Matrix.plan(shape, right) == Done{first} : Result<&2, &2, Matrix.Error, Matrix.Plan<shape, right>>}
357  for first_same: {Shape.dimensions(Matrix.output_shape(shape, right, first)) == Shape.dimensions(left) : List<&2, U32>}
358  for +second: Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>
359  for second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) == Done{second} :
360    Result<&2, &2, Matrix.Error, Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>>}
361  for second_same: {Shape.dimensions(Matrix.product_shape(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)) ==
362    Shape.dimensions(right) : List<&2, U32>}
363  for +fill: F32
364  for +ops: List<&2, StorageLaws.Op>
365  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
366  for +left_fill: F32
367  for +left_ops: List<&2, StorageLaws.Op>
368  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
369  for +right_fill: F32
370  for +right_ops: List<&2, StorageLaws.Op>
371  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
372  MatmulPairContract<a, left, Contracts.matrix_rows(shape, right, first), Contracts.matrix_columns(shape, right, first),
373    row => column => DotFold.dot_fold(Contracts.factors(row, column,
374      Contracts.matrix_inner(shape, right, first), Contracts.matrix_columns(shape, right, first),
375      shape, fill, ops, right, right_fill, right_ops, Contracts.matrix_inner(shape, right, first)), 0.0),
376    b, right,
377    Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
378    Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
379    row => column => DotFold.dot_fold(Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{},
380      Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
381      Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
382      Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
383      shape, fill, ops, left, left_fill, left_ops, Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)), 0.0),
384    Backward.vjp(Backward.MatmulNT{a, b, left, right},
385      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
386        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
387      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
388
389# What a VJP that sends one contribution of a matrix promises: it goes to node
390# b with target's dimensions and holds cell(row, column) at row * columns +
391# column for every row < rows and column < columns.
392type MatmulSoleContract<-b: Nat, -target: Shape.Shape, -rows: Nat, -columns: Nat, -cell: Nat -> Nat -> F32,
393  -pass: Backward.Pass> is Type:
394  MatmulSoleContract{
395    target_id: {TensorBackward.sole_target(TensorBackward.sole(pass)) == Some{b} : Maybe<&2, Nat>},
396    dimensions: {TensorBackward.sole_dimensions(TensorBackward.sole(pass)) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
397    value: @+row: Nat -> @row_below: {Nat.is_lt(row, rows) == True{} : Bool} ->
398      @+column: Nat -> @column_below: {Nat.is_lt(column, columns) == True{} : Bool} ->
399      {TensorBackward.sole_point(TensorBackward.sole(pass), U32.from_nat(Nat.add(Nat.mul(row, columns), column))) ==
400        Some{cell(row, column)} : Maybe<&2, F32>}}
401
402# d(x @ W^T)/dW of a product whose input x is a constant: the VJP of
403# Frozen{a, MatmulNT{a, b, left, right}}, for an input a that is not the
404# weight b, sends one contribution, to b, with right's dimensions: g^T @ x,
405# cell (j, l) the ordered dot product over i of g[i, j] * x[i, l] from +0.0,
406# the words of the second contribution of tensor_matmul_nt_vjp_values. It
407# never builds g @ W.
408law tensor_matmul_nt_frozen_vjp_values:
409  for +a: Nat
410  for +b: Nat
411  for apart: {Nat.is_eq(a, b) == False{} : Bool}
412  for +left: Shape.Shape
413  for +right: Shape.Shape
414  for +shape: Shape.Shape
415  for +second: Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>
416  for second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) == Done{second} :
417    Result<&2, &2, Matrix.Error, Matrix.Product<Matrix.Swapped{}, Matrix.Swapped{}, shape, left>>}
418  for second_same: {Shape.dimensions(Matrix.product_shape(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)) ==
419    Shape.dimensions(right) : List<&2, U32>}
420  for +fill: F32
421  for +ops: List<&2, StorageLaws.Op>
422  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
423  for +left_fill: F32
424  for +left_ops: List<&2, StorageLaws.Op>
425  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
426  for +right_fill: F32
427  for +right_ops: List<&2, StorageLaws.Op>
428  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
429  MatmulSoleContract<b, right,
430    Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
431    Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
432    row => column => DotFold.dot_fold(Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{},
433      Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
434      Contracts.product_rows(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
435      Contracts.product_columns(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second),
436      shape, fill, ops, left, left_fill, left_ops, Contracts.product_inner(Matrix.Swapped{}, Matrix.Swapped{}, shape, left, second)), 0.0),
437    Backward.vjp(Backward.Frozen{a, Backward.MatmulNT{a, b, left, right}},
438      Backward.Operands{Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
439        Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}},
440      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
441
442# matmul_nt reads what matmul reads from the transposed copy: for a right
443# operand of dimensions [columns, inner], the factor pairs of cell
444# (row, column) of a product of two `Kept` operands (`matmul_nt`), left[row,
445# k] and right[column, k], are the pairs `right_transposed_factors` reads
446# through `permute_source` on right^T, the pairs `tensor_matmul_vjp_values`
447# states for g @ y^T. Both fold the same words in the same order, so
448# `matmul_nt(g, y)` and `matmul(g, transpose_copy(y))` agree word for word.
449law nt_reads_transpose:
450  for +row: Nat
451  for +column: Nat
452  for +rows: Nat
453  for +left: Shape.Shape
454  for +left_fill: F32
455  for +left_ops: List<&2, StorageLaws.Op>
456  for +right: Shape.Shape
457  for +right_fill: F32
458  for +right_ops: List<&2, StorageLaws.Op>
459  for +columns: U32
460  for +inner: U32
461  for matrix: {Shape.dimensions(right) == [columns, inner] : List<&2, U32>}
462  for column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool}
463  for +count: Nat
464  for room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}
465  {Contracts.product_factors(row, column, Matrix.Kept{}, Matrix.Kept{}, U32.to_nat(inner), rows, U32.to_nat(columns),
466    left, left_fill, left_ops, right, right_fill, right_ops, count) ==
467    right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops,
468      right, right_fill, right_ops, swapped(right), count) : List<&2, DotFold.Factors>}
469
470# matmul_tn likewise: for a left operand of dimensions [inner, rows], the
471# pairs of a product of two `Swapped` operands (`matmul_tn`), left[k, row] and
472# right[k, column], are the pairs `left_transposed_factors` reads through
473# `permute_source` on left^T, the pairs `tensor_matmul_vjp_values` states for
474# x^T @ g.
475law tn_reads_transpose:
476  for +row: Nat
477  for +column: Nat
478  for +columns: Nat
479  for +left: Shape.Shape
480  for +left_fill: F32
481  for +left_ops: List<&2, StorageLaws.Op>
482  for +right: Shape.Shape
483  for +right_fill: F32
484  for +right_ops: List<&2, StorageLaws.Op>
485  for +inner: U32
486  for +rows: U32
487  for matrix: {Shape.dimensions(left) == [inner, rows] : List<&2, U32>}
488  for row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool}
489  for +count: Nat
490  for room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}
491  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(inner), U32.to_nat(rows), columns,
492    left, left_fill, left_ops, right, right_fill, right_ops, count) ==
493    left_transposed_factors(row, column, U32.to_nat(inner), columns, left, left_fill, left_ops, swapped(left),
494      right, right_fill, right_ops, count) : List<&2, DotFold.Factors>}
495
496# What a VJP that sends one contribution promises: it goes to node a with the
497# target's dimensions and holds cell(i) at every in-bounds i.
498type OneContract<-a: Nat, -target: Shape.Shape, -cell: U32 -> Maybe<&2, F32>, -pass: Backward.Pass> is Type:
499  OneContract{
500    target_id: {TensorBackward.sole_target(TensorBackward.sole(pass)) == Some{a} : Maybe<&2, Nat>},
501    dimensions: {TensorBackward.sole_dimensions(TensorBackward.sole(pass)) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
502    value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
503      {TensorBackward.sole_point(TensorBackward.sole(pass), index) == cell(index) : Maybe<&2, F32>}}
504
505# d relu(a)/da: with the saved input x (input's shape) and the gradient g
506# (shape), the Relu VJP sends the pointwise ReluBackward of x and g, reduced
507# to the input's shape, to a, and nothing else. Word k of the product is
508# scalar(ReluBackward, x[left_map(plan, k)], g[right_map(plan, k)]): +0.0
509# where that word of x is at most zero (`at_most_zero`, IEEE x <= 0), else
510# g's own word, as PyTorch's threshold_backward(g, relu(x), 0). The rank-0
511# value saved beside x is not read.
512law tensor_relu_vjp_values:
513  for +a: Nat
514  for +input: Shape.Shape
515  for +shape: Shape.Shape
516  for +plan: Shape.BroadcastPlan
517  for planned: {Shape.broadcast_checked(input, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
518  for +input_plan: Shape.BroadcastPlan
519  for input_planned: {Shape.broadcast_checked(input, Shape.output_shape(plan)) == Done{input_plan} :
520    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
521  for input_same: {Shape.dimensions(Shape.output_shape(input_plan)) == Shape.dimensions(Shape.output_shape(plan)) : List<&2, U32>}
522  for zero: Core.Dynamic
523  for +fill: F32
524  for +ops: List<&2, StorageLaws.Op>
525  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
526  for +input_fill: F32
527  for +input_ops: List<&2, StorageLaws.Op>
528  for -input_fits: {Lifecycle.logical_length(Replayed.storage(input, input_fill, input_ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}
529  OneContract<a, input, index => reduced(input, Shape.output_shape(plan), input_plan,
530      cells(Elementwise.ReluBackward{}, plan, input, input_fill, input_ops, shape, fill, ops,
531        Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))),
532      index),
533    Backward.vjp(Backward.Relu{a, input},
534      Backward.Operands{Core.Packed{input, Core.TTensor{Replayed.storage(input, input_fill, input_ops), input_fits}}, zero},
535      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
536
537# d mse_loss(a, b)/da and d/db: with the saved differences e = right - left
538# and d = left - right (both of the broadcast shape `diff`) and the gradient
539# g (shape), all replayed, the MseLoss VJP sends the pointwise MseGrad of g
540# and d, reduced to the left operand's shape, to a, then the pointwise
541# MseGrad of g and e, reduced to the right operand's shape, to b. Word k of
542# the first product is scalar(MseGrad{2 / N}, g[left_map(plan, k)],
543# d[right_map(plan, k)]) = (2 / N * d) * g, PyTorch's mse_loss_backward,
544# with N the cells of `diff` (Loss.mse_scale).
545law tensor_mse_vjp_values:
546  for +a: Nat
547  for +b: Nat
548  for +left: Shape.Shape
549  for +right: Shape.Shape
550  for +diff: Shape.Shape
551  for +shape: Shape.Shape
552  for +plan: Shape.BroadcastPlan
553  for planned: {Shape.broadcast_checked(shape, diff) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
554  for +left_plan: Shape.BroadcastPlan
555  for left_planned: {Shape.broadcast_checked(left, Shape.output_shape(plan)) == Done{left_plan} :
556    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
557  for left_same: {Shape.dimensions(Shape.output_shape(left_plan)) == Shape.dimensions(Shape.output_shape(plan)) : List<&2, U32>}
558  for +right_plan: Shape.BroadcastPlan
559  for right_planned: {Shape.broadcast_checked(right, Shape.output_shape(plan)) == Done{right_plan} :
560    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
561  for right_same: {Shape.dimensions(Shape.output_shape(right_plan)) == Shape.dimensions(Shape.output_shape(plan)) : List<&2, U32>}
562  for +fill: F32
563  for +ops: List<&2, StorageLaws.Op>
564  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
565  for +e_fill: F32
566  for +e_ops: List<&2, StorageLaws.Op>
567  for -e_fits: {Lifecycle.logical_length(Replayed.storage(diff, e_fill, e_ops)) == Shape.logical_numel(Shape.dimensions(diff)) : Nat}
568  for +d_fill: F32
569  for +d_ops: List<&2, StorageLaws.Op>
570  for -d_fits: {Lifecycle.logical_length(Replayed.storage(diff, d_fill, d_ops)) == Shape.logical_numel(Shape.dimensions(diff)) : Nat}
571  PairContract<a, left, index => reduced(left, Shape.output_shape(plan), left_plan,
572      cells(Elementwise.MseGrad{Loss.mse_scale(diff)}, plan, shape, fill, ops, diff, d_fill, d_ops,
573        Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))),
574      index),
575    b, right, index => reduced(right, Shape.output_shape(plan), right_plan,
576      cells(Elementwise.MseGrad{Loss.mse_scale(diff)}, plan, shape, fill, ops, diff, e_fill, e_ops,
577        Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))),
578      index),
579    Backward.vjp(Backward.MseLoss{a, b, left, right},
580      Backward.Operands{Core.Packed{diff, Core.TTensor{Replayed.storage(diff, e_fill, e_ops), e_fits}},
581        Core.Packed{diff, Core.TTensor{Replayed.storage(diff, d_fill, d_ops), d_fits}}},
582      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
583
584# d cross_entropy(a)/da: with the saved log_softmax lp and one-hot rows
585# (both of the input's shape) and the gradient g (shape), all replayed, the
586# CrossEntropy VJP sends two contributions, both to a with the input's
587# dimensions: the pointwise OneHotGrad of g and the one-hot rows, then the
588# pointwise SoftmaxGrad of g and lp, each reduced to the input's shape. Word
589# k of the first is -(g / N) where the one-hot word is 1.0 and +0.0
590# elsewhere; word k of the second is -(exp(lp) * r), with r = -(g / N) when
591# the input has one column (Loss.one_class) and r = -(g / N) + 0.0 otherwise;
592# N is the input's first dimension as an F32 (Loss.count_f32(Loss.rows(input))).
593# Both cases are stated: SoftmaxGrad carries Loss.one_class(input).
594# The slot adds the second into the first, PyTorch's one_hot gradient minus
595# exp(lp) times its row sum.
596law tensor_cross_entropy_vjp_values:
597  for +a: Nat
598  for +input: Shape.Shape
599  for +shape: Shape.Shape
600  for +plan: Shape.BroadcastPlan
601  for planned: {Shape.broadcast_checked(shape, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
602  for +input_plan: Shape.BroadcastPlan
603  for input_planned: {Shape.broadcast_checked(input, Shape.output_shape(plan)) == Done{input_plan} :
604    Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
605  for input_same: {Shape.dimensions(Shape.output_shape(input_plan)) == Shape.dimensions(Shape.output_shape(plan)) : List<&2, U32>}
606  for +fill: F32
607  for +ops: List<&2, StorageLaws.Op>
608  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
609  for +lp_fill: F32
610  for +lp_ops: List<&2, StorageLaws.Op>
611  for -lp_fits: {Lifecycle.logical_length(Replayed.storage(input, lp_fill, lp_ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}
612  for +hot_fill: F32
613  for +hot_ops: List<&2, StorageLaws.Op>
614  for -hot_fits: {Lifecycle.logical_length(Replayed.storage(input, hot_fill, hot_ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}
615  PairContract<a, input, index => reduced(input, Shape.output_shape(plan), input_plan,
616      cells(Elementwise.OneHotGrad{Loss.count_f32(Loss.rows(input))}, plan, shape, fill, ops, input, hot_fill, hot_ops,
617        Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))),
618      index),
619    a, input, index => reduced(input, Shape.output_shape(plan), input_plan,
620      cells(Elementwise.SoftmaxGrad{Loss.count_f32(Loss.rows(input)), Loss.one_class(input)}, plan, shape, fill, ops, input, lp_fill, lp_ops,
621        Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))),
622      index),
623    Backward.vjp(Backward.CrossEntropy{a, input},
624      Backward.Operands{Core.Packed{input, Core.TTensor{Replayed.storage(input, lp_fill, lp_ops), lp_fits}},
625        Core.Packed{input, Core.TTensor{Replayed.storage(input, hot_fill, hot_ops), hot_fits}}},
626      Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>