1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/LAWS.bend as StorageLaws 4 import ../tensor/core.bend as Core 5 import ../tensor/layout.bend as Layout 6 import ../tensor/lifecycle.bend as Lifecycle 7 import ../tensor/elementwise.bend as Elementwise 8 import ../tensor/reduction.bend as Reduction 9 import ../tensor/matmul.bend as Matrix 10 import ../tensor/dot_fold.bend as DotFold 11 import ../tensor/transpose_spec.bend as Transpose 12 import ../tensor/observation.bend as Values 13 import ../tensor/contracts/replayed.bend as Replayed 14 import ../tensor/contracts/LAWS.bend as Contracts 15 import ../tensor/contracts/binary_LAWS.bend as Binary 16 import ./sum_to_size_fold_LAWS.bend as SumToSizeFold 17 import ./tensor_backward.bend as Backward 18 import ./loss.bend as Loss 19 import ./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). 30 def 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. 38 def 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`. 46 def 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. 60 law 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. 78 def 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 92 def 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). 109 type 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. 127 law 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 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. 152 law 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 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)]). 183 law 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 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. 223 def swapped(+shape: Shape.Shape) -> List<&2, U32>: 224 Transpose.swapped_axes(shape, 0, 1) 225 226 def transposed(+shape: Shape.Shape, +plan: Layout.Plan) -> 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. 233 def 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. 251 def 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. 268 type 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. 292 law 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 300 for right_laid: {Layout.permutation(right, swapped(right)) == Done{right_plan} : Result<&2, &2, Layout.Error, Layout.Plan>} 301 for +first: Matrix.Plan 302 for first_planned: {Matrix.plan(shape, transposed(right, right_plan)) == Done{first} : 303 Result<&2, &2, Matrix.Error, Matrix.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 307 for left_laid: {Layout.permutation(left, swapped(left)) == Done{left_plan} : Result<&2, &2, Layout.Error, Layout.Plan>} 308 for +second: Matrix.Plan 309 for second_planned: {Matrix.plan(transposed(left, left_plan), shape) == Done{second} : 310 Result<&2, &2, Matrix.Error, Matrix.Plan>} 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 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. 349 law 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 356 for first_planned: {Matrix.plan(shape, right) == Done{first} : Result<&2, &2, Matrix.Error, Matrix.Plan>} 357 for first_same: {Shape.dimensions(Matrix.output_shape(shape, right, first)) == Shape.dimensions(left) : List<&2, U32>} 358 for +second: Matrix.Product 359 for second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) == Done{second} : 360 Result<&2, &2, Matrix.Error, Matrix.Product>} 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 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. 392 type 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. 408 law 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 416 for second_planned: {Matrix.product_plan(Matrix.Swapped{}, Matrix.Swapped{}, shape, left) == Done{second} : 417 Result<&2, &2, Matrix.Error, Matrix.Product>} 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 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. 449 law 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. 475 law 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. 498 type 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. 512 law 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 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). 545 law 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 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. 596 law 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 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}})>