src/autograd/tensor_backward_LAWS.bend

bendtorch-core@0.1.0.0 · 418 lines · 37 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
6import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
7import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
8import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
9import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
10import bendtorch-core@0.1.0.0/src/autograd/sum_dims.bend as SumDims
11import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
12import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
13
14# The nodes of a list of entries, newest first.
15def nodes_of(entries: List<&1, Backward.Entry>) -> List<&2, Backward.Node>:
16  match entries:
17    case Nil{}:
18      Nil{}
19    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
20      node <> nodes_of(rest)
21
22def slots_of(entries: List<&1, Backward.Entry>) -> List<&1, Gradient.Slot>:
23  match entries:
24    case Nil{}:
25      Nil{}
26    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
27      slot <> slots_of(rest)
28
29def tape_nodes(tape: Backward.Tape) -> List<&2, Backward.Node>:
30  Backward.Tape{entries} = tape
31  nodes_of(entries)
32
33def tape_slots(tape: Backward.Tape) -> List<&1, Gradient.Slot>:
34  Backward.Tape{entries} = tape
35  slots_of(entries)
36
37# Operand bounds, stated without the engine's validator: a node may name only
38# nodes older than itself, so each operand id must be below the number of
39# older nodes.
40def targets(node: Backward.Node) -> List<&2, Nat>:
41  match node:
42    case Backward.Leaf{}:
43      []
44    case Backward.Add{a, b, left, right}:
45      [a, b]
46    case Backward.Sub{a, b, left, right}:
47      [a, b]
48    case Backward.Mul{a, b, left, right}:
49      [a, b]
50    case Backward.Sum{a, input}:
51      [a]
52    case Backward.SumDims{a, input, kept}:
53      [a]
54    case Backward.Matmul{a, b, left, right}:
55      [a, b]
56    case Backward.MatmulNT{a, b, left, right}:
57      [a, b]
58    case Backward.Relu{a, input}:
59      [a]
60    case Backward.MseLoss{a, b, left, right}:
61      [a, b]
62    case Backward.CrossEntropy{a, input}:
63      [a]
64    case Backward.BatchedMatmul{a, b, left, right}:
65      [a, b]
66    case Backward.Slice{a, input, kept, plan}:
67      [a]
68    case Backward.Gathered{a, input, plan, entries}:
69      [a]
70    case Backward.Constant{}:
71      []
72    case Backward.Frozen{id, inner}:
73      targets(inner)
74
75def bounded(ids: List<&2, Nat>, +bound: Nat) -> Bool:
76  match ids:
77    case Nil{}:
78      True{}
79    case Con{+id, rest}:
80      Bool.and(Nat.is_lt(id, bound), bounded(rest, bound))
81
82def well_formed(nodes: List<&2, Backward.Node>) -> Bool:
83  match nodes:
84    case Nil{}:
85      True{}
86    case Con{node, +rest}:
87      Bool.and(bounded(targets(node), List.length(&2, Backward.Node, rest)), well_formed(rest))
88
89# The slots newer than a target, put back in front of the result.
90def relisted(newer: List<&1, Gradient.Slot>, result: Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>) ->
91  Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>:
92  match newer:
93    case Nil{}:
94      result
95    case Con{slot, rest}:
96      Backward.prepend(Gradient.Slot, slot, relisted(rest, result))
97
98# The entries newer than the output, each with no gradient, in front of the
99# sweep's result.
100def newer_unset(newer: List<&1, Backward.Entry>, swept: Backward.Swept()) -> Backward.Swept():
101  match newer:
102    case Nil{}:
103      swept
104    case Con{entry, rest}:
105      Backward.finished(entry, Gradient.Unset{}, newer_unset(rest, swept))
106
107# The only contribution of a VJP, when it sends exactly one.
108def sole(pass: Backward.Pass) -> Maybe<&1, Backward.Contribution>:
109  Backward.Pass{saved, g, contributions} = pass
110  match contributions:
111    case Fail{error}:
112      None{}
113    case Done{Nil{}}:
114      None{}
115    case Done{Con{contribution, Nil{}}}:
116      Some{contribution}
117    case Done{Con{contribution, Con{other, rest}}}:
118      None{}
119
120def sole_target(found: Maybe<&1, Backward.Contribution>) -> Maybe<&2, Nat>:
121  match found:
122    case None{}:
123      None{}
124    case Some{Backward.Contribution{target, shape, gradient}}:
125      Some{target}
126
127def sole_dimensions(found: Maybe<&1, Backward.Contribution>) -> Maybe<&2, List<&2, U32>>:
128  match found:
129    case None{}:
130      None{}
131    case Some{Backward.Contribution{target, shape, Core.Packed{dimensions, tensor}}}:
132      Some{Shape.dimensions(dimensions)}
133
134def sole_point(found: Maybe<&1, Backward.Contribution>, +index: U32) -> Maybe<&2, F32>:
135  match found:
136    case None{}:
137      None{}
138    case Some{Backward.Contribution{target, shape, gradient}}:
139      Core.dynamic_read_value(Core.read_dynamic(gradient, index))
140
141# What a VJP that sends one contribution promises: it goes to node a, it has
142# the target's dimensions, and every in-bounds index i holds cell(i).
143type SoleContract<-a: Nat, -target: Shape.Shape, -cell: Nat -> F32, -pass: Backward.Pass> is Type:
144  SoleContract{
145    target_id: {sole_target(sole(pass)) == Some{a} : Maybe<&2, Nat>},
146    dimensions: {sole_dimensions(sole(pass)) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
147    value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
148      {sole_point(sole(pass), index) == Some{cell(U32.to_nat(index))} : Maybe<&2, F32>}}
149
150# NegateLeft flips the sign bit of the left word, as PyTorch's neg does, and
151# ignores the right word. The Sub VJP sends -g to its second operand with it.
152law negate_left_scalar:
153  for +left: F32
154  for +right: F32
155  {Elementwise.scalar(Elementwise.NegateLeft{}, left, right) == F32.neg(left) : F32}
156
157# A leaf sends nothing, and hands back its saved values and its gradient.
158law tensor_vjp_leaf:
159  for saved: Backward.Saved
160  for g: Core.Dynamic
161  {Backward.vjp(Backward.Leaf{}, saved, g) == Backward.Pass{saved, g, Done{Nil{}}} : Backward.Pass}
162
163# A constant sends nothing, and hands back its saved values and its gradient.
164law tensor_vjp_constant:
165  for saved: Backward.Saved
166  for g: Core.Dynamic
167  {Backward.vjp(Backward.Constant{}, saved, g) == Backward.Pass{saved, g, Done{Nil{}}} : Backward.Pass}
168
169# Frozen{id, node}, the node of a recording whose operand id is a constant,
170# sends node's contributions without the ones to id, in order, and hands back
171# what node's VJP hands back. Where Backward.skips(id, node) holds (node is
172# x @ W^T and id is x and not W), the VJP does not compute the product it
173# would drop; tensor_matmul_nt_frozen_vjp_values states what it sends.
174law tensor_vjp_frozen_unsent:
175  for +id: Nat
176  for +node: Backward.Node
177  for saved: Backward.Saved
178  for g: Core.Dynamic
179  for computed: {Backward.skips(id, node) == False{} : Bool}
180  {Backward.vjp(Backward.Frozen{id, node}, saved, g) == Backward.unsent(id, Backward.vjp(node, saved, g)) : Backward.Pass}
181
182# Whether a list of contributions holds one to id.
183def reaches(+id: Nat, items: List<&1, Backward.Contribution>) -> Bool:
184  match items:
185    case Nil{}:
186      False{}
187    case Con{Backward.Contribution{target, shape, gradient}, rest}:
188      Bool.or(Nat.is_eq(id, target), reaches(id, rest))
189
190# Whether a VJP's pass sends a contribution to id.
191def sends_to(+id: Nat, pass: Backward.Pass) -> Bool:
192  Backward.Pass{saved, g, contributions} = pass
193  match contributions:
194    case Fail{error}:
195      False{}
196    case Done{items}:
197      reaches(id, items)
198
199# A frozen operand gets nothing: for any node, saved values and g, no
200# contribution of the VJP of Frozen{id, node} goes to id, skipped product
201# included.
202law tensor_vjp_frozen_spares:
203  for +id: Nat
204  for +node: Backward.Node
205  for saved: Backward.Saved
206  for g: Core.Dynamic
207  {sends_to(id, Backward.vjp(Backward.Frozen{id, node}, saved, g)) == False{} : Bool}
208
209# backward of an output the tape does not hold is refused with UnknownNode,
210# and the tape comes back unchanged.
211law tensor_backward_unknown:
212  for entries: List<&1, Backward.Entry>
213  for +output: Nat
214  for outside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == False{} : Bool}
215  {Backward.backward(Backward.Tape{entries}, output) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
216    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
217
218# A tape with a node that refers to itself, to a newer node or to a missing
219# node is refused with DanglingOperand, and comes back unchanged.
220law tensor_backward_dangling:
221  for entries: List<&1, Backward.Entry>
222  for +output: Nat
223  for inside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
224  for unsound: {Backward.wired(nodes_of(entries), List.length(&2, Backward.Node, nodes_of(entries))) == False{} : Bool}
225  {Backward.backward(Backward.Tape{entries}, output) == (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) :
226    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
227
228# Otherwise backward runs pass from the selected output and accumulates its
229# working gradients into the slots.
230law tensor_backward_success:
231  for entries: List<&1, Backward.Entry>
232  for +output: Nat
233  for inside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
234  for sound: {Backward.wired(nodes_of(entries), List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
235  {Backward.backward(Backward.Tape{entries}, output) ==
236    Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, nodes_of(entries)), output,
237      Backward.at_head(List.length(&2, Backward.Node, nodes_of(entries)), output))) :
238    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
239
240# The engine's wiring check is exactly the operand-bound specification above.
241law tensor_wired_matches_operand_bounds:
242  for +graph: List<&2, Backward.Node>
243  {Backward.wired(graph, List.length(&2, Backward.Node, graph)) == well_formed(graph) : Bool}
244
245# The seed is the selected output (C1): for any entries newer than it, those
246# entries get no gradient, and the sweep starts at the output with ones of the
247# output's shape, Core.new(shape, 1.0), and no gradient for the older entries.
248law tensor_seed_selected_output:
249  for newer: List<&1, Backward.Entry>
250  for +node: Backward.Node
251  for +shape: Shape.Shape
252  for value: Core.Dynamic
253  for saved: Backward.Saved
254  for slot: Gradient.Slot
255  for rest: List<&1, Backward.Entry>
256  for +older: Nat
257  for fits: {List.length(&1, Backward.Entry, rest) == older : Nat}
258  {Backward.pass(List.append(&1, Backward.Entry, newer, Backward.Entry{node, shape, value, saved, slot} <> rest),
259      1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older,
260      Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older)) ==
261    newer_unset(newer, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0))) : Backward.Swept()}
262
263# The sweep starts from the seed: the output gets ones of its shape and every
264# older entry starts with no gradient.
265law tensor_seeded_sweep:
266  for +node: Backward.Node
267  for +shape: Shape.Shape
268  for value: Core.Dynamic
269  for saved: Backward.Saved
270  for slot: Gradient.Slot
271  for rest: List<&1, Backward.Entry>
272  for +below: Nat
273  for ones: Core.Tensor<shape>
274  {Backward.seeded(node, shape, value, saved, slot, rest, below, Done{ones}) ==
275    Backward.sweep(Backward.Entry{node, shape, value, saved, slot} <> rest,
276      Done{Gradient.Held{Core.Packed{shape, ones}} <> Backward.unsets(below)}, 1n+below) : Backward.Swept()}
277
278# An entry that received no gradient sends nothing: the older entries are
279# swept with their working gradients unchanged.
280law tensor_unreached_sends_nothing:
281  for +node: Backward.Node
282  for +shape: Shape.Shape
283  for value: Core.Dynamic
284  for saved: Backward.Saved
285  for slot: Gradient.Slot
286  for rest: List<&1, Backward.Entry>
287  for pending: List<&1, Gradient.Slot>
288  for +below: Nat
289  {Backward.sweep(Backward.Entry{node, shape, value, saved, slot} <> rest, Done{Gradient.Unset{} <> pending}, 1n+below) ==
290    Backward.finished(Backward.Entry{node, shape, value, saved, slot}, Gradient.Unset{},
291      Backward.sweep(rest, Done{pending}, below)) : Backward.Swept()}
292
293# The contributions of one VJP are delivered in list order, the first operand
294# before the second.
295law tensor_deliver_in_order:
296  for +target: Nat
297  for +shape: Shape.Shape
298  for gradient: Core.Dynamic
299  for rest: List<&1, Backward.Contribution>
300  for +count: Nat
301  for slots: List<&1, Gradient.Slot>
302  {Backward.deliver(Backward.Contribution{target, shape, gradient} <> rest, count, Done{slots}) ==
303    Backward.deliver(rest, count, Backward.bump(slots, count, target, shape, gradient, Backward.at_head(count, target))) :
304    Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
305
306# A contribution to node target reaches exactly the slot of that node, whose
307# Gradient.accumulate adds it to the held total (or moves it in); every other
308# slot is kept, whatever the list holds.
309law tensor_bump_target:
310  for newer: List<&1, Gradient.Slot>
311  for slot: Gradient.Slot
312  for rest: List<&1, Gradient.Slot>
313  for +shape: Shape.Shape
314  for gradient: Core.Dynamic
315  for +older: Nat
316  for fits: {List.length(&1, Gradient.Slot, rest) == older : Nat}
317  {Backward.bump(List.append(&1, Gradient.Slot, newer, slot <> rest), 1n+Nat.add(older, List.length(&1, Gradient.Slot, newer)),
318      older, shape, gradient, Backward.at_head(1n+Nat.add(older, List.length(&1, Gradient.Slot, newer)), older)) ==
319    relisted(newer, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient))) :
320    Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
321
322# After the sweep, an entry that received a gradient g accumulates it into its
323# slot with Gradient.accumulate, so a repeated backward adds; an entry that
324# received nothing keeps its slot.
325law tensor_merge_accumulates:
326  for +node: Backward.Node
327  for +shape: Shape.Shape
328  for value: Core.Dynamic
329  for saved: Backward.Saved
330  for slot: Gradient.Slot
331  for rest: List<&1, Backward.Entry>
332  for g: Core.Dynamic
333  for more: List<&1, Gradient.Slot>
334  {Backward.merge(Backward.Entry{node, shape, value, saved, slot} <> rest, Gradient.Held{g} <> more) ==
335    Backward.merged_entry(node, shape, value, saved, Gradient.accumulate(shape, slot, g), Backward.merge(rest, more)) :
336    List<&1, Backward.Entry> & Result<&2, &2, Backward.Error, Unit>}
337
338law tensor_merge_keeps_unreached:
339  for +node: Backward.Node
340  for +shape: Shape.Shape
341  for value: Core.Dynamic
342  for saved: Backward.Saved
343  for slot: Gradient.Slot
344  for rest: List<&1, Backward.Entry>
345  for more: List<&1, Gradient.Slot>
346  {Backward.merge(Backward.Entry{node, shape, value, saved, slot} <> rest, Gradient.Unset{} <> more) ==
347    Backward.merged_entry(node, shape, value, saved, Gradient.Accumulated{slot}, Backward.merge(rest, more)) :
348    List<&1, Backward.Entry> & Result<&2, &2, Backward.Error, Unit>}
349
350# zero_grad empties every slot, like PyTorch's zero_grad(set_to_none=True),
351# and keeps every node.
352law tensor_zero_grad_empties:
353  for entries: List<&1, Backward.Entry>
354  {tape_slots(Backward.zero_grad(Backward.Tape{entries})) == Backward.unsets(List.length(&1, Backward.Entry, entries)) :
355    List<&1, Gradient.Slot>}
356
357law tensor_zero_grad_keeps_nodes:
358  for entries: List<&1, Backward.Entry>
359  {tape_nodes(Backward.zero_grad(Backward.Tape{entries})) == nodes_of(entries) : List<&2, Backward.Node>}
360
361# d(sum a)/da through the live broadcast_to_pointwise law: for a replayed
362# gradient g whose shape broadcasts to exactly a's shape, the Sum VJP sends one
363# contribution, to a, with a's dimensions, and every position i holds the
364# gradient word at left_map(plan, i), its right-aligned broadcast projection;
365# for the rank-0 g of a sum, that is g's only word at every position.
366law tensor_sum_vjp_pointwise:
367  for +a: Nat
368  for +shape: Shape.Shape
369  for +input: Shape.Shape
370  for +plan: Shape.BroadcastPlan
371  for planned: {Shape.broadcast_checked(shape, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
372  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>}
373  for saved: Backward.Saved
374  for +fill: F32
375  for +ops: List<&2, StorageLaws.Op>
376  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
377  SoleContract<a, input, index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
378    Backward.vjp(Backward.Sum{a, input}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
379
380# d(sum of a over some axes)/da: a gradient g of `shape`, the kept or the
381# dropped shape, with as many elements as the kept shape, is read with the
382# kept shape and broadcast as the Sum VJP broadcasts: for a replayed g whose
383# words are a replay over the kept shape's count, and a kept shape that
384# broadcasts to exactly a's shape, the VJP sends one contribution, to a, with
385# a's dimensions, and every position i holds the word of g at
386# left_map(plan, i), the kept shape's right-aligned projection of i.
387law tensor_sum_dims_vjp_pointwise:
388  for +a: Nat
389  for +shape: Shape.Shape
390  for +input: Shape.Shape
391  for +kept: Shape.Shape
392  for +plan: Shape.BroadcastPlan
393  for planned: {Shape.broadcast_checked(kept, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
394  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>}
395  for counted: {Nat.is_eq(Shape.elements(shape), Shape.elements(kept)) == True{} : Bool}
396  for saved: Backward.Saved
397  for +fill: F32
398  for +ops: List<&2, StorageLaws.Op>
399  for -fits: {Lifecycle.logical_length(Replayed.storage(kept, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
400  SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
401    Backward.vjp(Backward.SumDims{a, input, kept}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}})>
402
403# The tape's sum over axes takes PyTorch's signed axes: they are wrapped
404# among the operand's axes, its rank or 1 at rank 0 (SumDims.limit), before
405# the plan checks them. So an axis outside -limit..limit-1 is refused with
406# AxisOutOfBounds and two axes that wrap to one with DuplicateAxis
407# (sum_dims_LAWS.bend, signed_LAWS.bend signed_axes_bounded).
408law tensor_sum_dims_wraps:
409  for +a: Nat
410  for +axes: List<&2, Signed.Position>
411  for keepdim: Bool
412  for entries: List<&1, Backward.Entry>
413  for input: Core.Dynamic
414  for +shape: Shape.Shape
415  {Backward.dims_shaped(a, axes, keepdim, entries, (input, shape)) ==
416    Backward.dims_planned(a, entries, input, shape, keepdim,
417      SumDims.plan(shape, Signed.wrap_all(SumDims.limit(Shape.dimensions(shape)), axes))) :
418    Backward.Tape & Result<&2, &2, Backward.Error, Nat>}