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/lifecycle.bend as Lifecycle 6 import ../tensor/elementwise.bend as Elementwise 7 import ../tensor/reduction.bend as Reduction 8 import ../tensor/word_conversion.bend as Conversion 9 import ../tensor/contracts/replayed.bend as Replayed 10 import ../tensor/contracts/pointwise.bend as Pointwise 11 import ../tensor/contracts/pointwise_proof.bend as PointwiseProof 12 import ../tensor/contracts/broadcast.bend as PointwiseBroadcast 13 import ../tensor/contracts/LAWS.bend as Contracts 14 import ../tensor/contracts/binary_proofs.bend as BinaryProof 15 # Proves binary_replayed, which importing the proof above would otherwise leave open. 16 import ../tensor/contracts/binary_PROOF.bend as BinaryLaw 17 import ./reduce_to_fold_PROOF.bend as ReduceTo 18 import ./sum_to_size_fold_LAWS.bend as SumToSizeFold 19 import ./sum_to_size_fold_PROOF.bend as SumToSizeFoldProof 20 import ./tensor_backward.bend as Backward 21 import ./loss.bend as Loss 22 import ./tensor_vjp_LAWS.bend as Laws 23 import ../tensor/layout.bend as Layout 24 import ../tensor/matmul.bend as Matrix 25 import ../tensor/transpose_spec.bend as Transpose 26 import ../tensor/dot_fold.bend as DotFold 27 import ../tensor/observation.bend as Values 28 import ../tensor/index_map.bend as IndexMap 29 import ../tensor/permute_values.bend as PermuteValues 30 import ../tensor/permute_values_records.bend as PermuteRecords 31 import ../tensor/permute_values_contract_LAWS.bend as PermuteLaws 32 import ../tensor/permute_values_contract_PROOF.bend as PermuteProof 33 import ../tensor/contracts/matmul_proof.bend as MatmulProof 34 import ../tensor/contracts/product_proof.bend as ProductProof 35 import ../tensor/transpose_PROOF.bend as TransposeProof 36 import ../tensor/matrix_index_proof.bend as MatrixIndex 37 38 def 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 44 def 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. 63 def 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 92 def 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 104 def 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. 109 def 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: 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 148 def 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 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 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 169 def 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 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 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 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 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 213 def 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 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 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 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 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. 267 def permuted_at(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan, 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 289 def 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. 304 def 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, +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. 353 def 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, +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. 402 def MatmulGoal(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape, 403 +right_plan: Layout.Plan, +first: Matrix.Plan, 404 +left_plan: Layout.Plan, +second: Matrix.Plan, 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 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. 422 def 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 446 def facts_left(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32, 447 facts: MatmulProof.MatrixFacts) -> {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}: 448 MatmulProof.MatrixFacts{le, re, oe, compatible} = facts 449 le 450 451 def facts_right(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32, 452 facts: MatmulProof.MatrixFacts) -> 453 {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}: 454 MatmulProof.MatrixFacts{le, re, oe, compatible} = facts 455 re 456 457 def facts_compatible(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32, 458 facts: MatmulProof.MatrixFacts) -> {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}: 459 MatmulProof.MatrixFacts{le, re, oe, compatible} = facts 460 compatible 461 462 def 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, 465 +right_laid: {Layout.permutation(right, Laws.swapped(right)) == Done{right_plan} : Result<&2, &2, Layout.Error, Layout.Plan>}, 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>}, 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, 477 +left_laid: {Layout.permutation(left, Laws.swapped(left)) == Done{left_plan} : Result<&2, &2, Layout.Error, Layout.Plan>}, 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>}, 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 = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible} 503 +second : Matrix.Plan = 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 = 506 MatmulProof.plan_facts(shape, flipped, first, first_planned) 507 +back_facts : MatmulProof.MatrixFacts = 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 581 def 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. 596 def 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. 622 def 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. 654 def NTGoal(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape, 655 +first: Matrix.Plan, +second: Matrix.Product, 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 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 672 def 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>}, 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>}, 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 = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible} 704 +second : Matrix.Product = 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 = 707 MatmulProof.plan_facts(shape, right, first, first_planned) 708 +back_facts : ProductProof.ProductFacts = 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 782 def 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. 794 def 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 820 def FrozenGoal(+b: Nat, +left: Shape.Shape, +right: Shape.Shape, +shape: Shape.Shape, 821 +second: Matrix.Product, 822 +fill: F32, +ops: List<&2, StorageLaws.Op>, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, pass: Backward.Pass) -> Type: 823 Laws.MatmulSoleContract 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 833 def 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>}, 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 = 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 = 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 898 def 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. 906 def 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 916 def 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]. 931 def 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. 939 def 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. 964 def 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 987 def 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 1018 def 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 1049 def 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 1054 def 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. 1061 def 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: 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 1082 def 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 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 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. 1104 def 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 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 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 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 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. 1157 def 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 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 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 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 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)