src/autograd/sgd_LAWS.bend

bendtorch-core@0.1.0.0 · 228 lines · 20 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/LAWS.bend as Contracts
9import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_LAWS.bend as Binary
10import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
11import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
12import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as TensorBackward
13import bendtorch-core@0.1.0.0/src/autograd/sgd.bend as Sgd
14
15# The cell of an SGD step: p - lr * g, with the product rounded before the
16# difference, F32 operation by F32 operation.
17law sgd_descend_scalar:
18  for +lr: F32
19  for +p: F32
20  for +g: F32
21  {Elementwise.scalar(Elementwise.Descend{lr}, p, g) == F32.sub(p, F32.mul(lr, g)) : F32}
22
23# A parameter whose slot is empty, PyTorch's p.grad is None, is left as it
24# was.
25law sgd_empty_slot_kept:
26  for +lr: F32
27  for node: Backward.Node
28  for +shape: Shape.Shape
29  for value: Core.Dynamic
30  for saved: Backward.Saved
31  {Sgd.descend(lr, Backward.Entry{node, shape, value, saved, Gradient.Unset{}}) ==
32    (Backward.Entry{node, shape, value, saved, Gradient.Unset{}}, Done{Unit{}}) : Backward.Entry & Result<&2, &2, Sgd.Error, Unit>}
33
34# What the transport form of sgd_step_values asks of P: P holds for the entry
35# with its node, shape and saved values, whose value is replayed over +0.0
36# with any history whose words are the cells of Descend{lr} over the old
37# value p and the gradient g, and whose slot still holds g with its words,
38# next to Done.
39def DescendWith(+lr: F32, node: Backward.Node, +shape: Shape.Shape, saved: Backward.Saved, +plan: Shape.BroadcastPlan,
40  +fill: F32, +ops: List<&2, StorageLaws.Op>,
41  +g_shape: Shape.Shape, +g_fill: F32, +g_ops: List<&2, StorageLaws.Op>,
42  -P: Backward.Entry & Result<&2, &2, Sgd.Error, Unit> -> Type) -> Type:
43  @+kept: List<&2, StorageLaws.Op> ->
44  @kept_words: {Contracts.words(shape, fill, kept, Shape.logical_numel(Shape.dimensions(shape))) ==
45    Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))) : List<&2, F32>} ->
46  @-kept_fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, kept)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
47  @+g_kept: List<&2, StorageLaws.Op> ->
48  @g_words: {Contracts.words(g_shape, g_fill, g_kept, Shape.logical_numel(Shape.dimensions(g_shape))) ==
49    Contracts.words(g_shape, g_fill, g_ops, Shape.logical_numel(Shape.dimensions(g_shape))) : List<&2, F32>} ->
50  @-g_fits: {Lifecycle.logical_length(Replayed.storage(g_shape, g_fill, g_kept)) == Shape.logical_numel(Shape.dimensions(g_shape)) : Nat} ->
51  @+output_ops: List<&2, StorageLaws.Op> ->
52  @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) ==
53    Binary.cells(Elementwise.Descend{lr}, plan, shape, fill, ops, g_shape, g_fill, g_ops,
54      Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} ->
55  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) ==
56    Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} ->
57    P((Backward.Entry{node, shape,
58      Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}},
59      saved, Gradient.Held{Core.Packed{g_shape, Core.TTensor{Replayed.storage(g_shape, g_fill, g_kept), g_fits}}}}, Done{Unit{}}))
60
61# A parameter whose slot holds g gets p - lr * g on the parallel elementwise
62# build, in transport form: for a replayed value p and a replayed gradient g
63# whose shapes broadcast, for every P, if P holds for every result that keeps
64# the node, the shape, the saved values and g with its words, and holds a
65# value of the plan's shape whose word k is scalar(Descend{lr},
66# p[left_map(plan, k)], g[right_map(plan, k)]) (sgd_descend_scalar: p - lr *
67# g), next to Done, then P holds for the step of that parameter.
68law sgd_step_values:
69  for +lr: F32
70  for node: Backward.Node
71  for +shape: Shape.Shape
72  for saved: Backward.Saved
73  for +g_shape: Shape.Shape
74  for +plan: Shape.BroadcastPlan
75  for planned: {Shape.broadcast_checked(shape, g_shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
76  for +fill: F32
77  for +ops: List<&2, StorageLaws.Op>
78  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
79  for +g_fill: F32
80  for +g_ops: List<&2, StorageLaws.Op>
81  for -g_fits: {Lifecycle.logical_length(Replayed.storage(g_shape, g_fill, g_ops)) == Shape.logical_numel(Shape.dimensions(g_shape)) : Nat}
82  for -P: Backward.Entry & Result<&2, &2, Sgd.Error, Unit> -> Type
83  for h: DescendWith(lr, node, shape, saved, plan, fill, ops, g_shape, g_fill, g_ops, P)
84  P(Sgd.descend(lr, Backward.Entry{node, shape, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, saved,
85    Gradient.Held{Core.Packed{g_shape, Core.TTensor{Replayed.storage(g_shape, g_fill, g_ops), g_fits}}}}))
86
87# The entries newer than the parameter, kept in front of the result.
88def newer_kept(newer: List<&1, Backward.Entry>, walked: Sgd.Walked()) -> Sgd.Walked():
89  match newer:
90    case Nil{}:
91      walked
92    case Con{entry, rest}:
93      Sgd.kept(entry, newer_kept(rest, walked))
94
95# The update of parameter id reaches exactly the entry with id older entries
96# below it, and keeps every other entry, whatever the list holds.
97law sgd_update_reaches_target:
98  for +lr: F32
99  for newer: List<&1, Backward.Entry>
100  for entry: Backward.Entry
101  for rest: List<&1, Backward.Entry>
102  for +older: Nat
103  for fits: {List.length(&1, Backward.Entry, rest) == older : Nat}
104  {Sgd.descend_at(lr, List.append(&1, Backward.Entry, newer, entry <> rest), 1n+Nat.add(older, List.length(&1, Backward.Entry, newer)),
105      older, Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older)) ==
106    newer_kept(newer, Sgd.placed(rest, Sgd.descend(lr, entry))) : Sgd.Walked()}
107
108# The parameters are updated in list order: none leaves the entries as they
109# are, and the first is updated before the rest, which see its result.
110law sgd_no_params:
111  for +lr: F32
112  for +count: Nat
113  for entries: List<&1, Backward.Entry>
114  {Sgd.descend_all(lr, count, Nil{}, entries) == (entries, Done{Unit{}}) : Sgd.Walked()}
115
116law sgd_params_in_order:
117  for +lr: F32
118  for +count: Nat
119  for +id: Nat
120  for rest: List<&2, Nat>
121  for entries: List<&1, Backward.Entry>
122  {Sgd.descend_all(lr, count, id <> rest, entries) ==
123    Sgd.continued(after => Sgd.descend_all(lr, count, rest, after), Sgd.descend_at(lr, entries, count, id, Backward.at_head(count, id))) :
124    Sgd.Walked()}
125
126# Whether parameter id is refused is read from the node with id older nodes
127# below it: a leaf is accepted, any other node refused with NotLeaf.
128law sgd_refusal_reads_target:
129  for newer: List<&2, Backward.Node>
130  for node: Backward.Node
131  for rest: List<&2, Backward.Node>
132  for +older: Nat
133  for fits: {List.length(&2, Backward.Node, rest) == older : Nat}
134  {Sgd.refusal_at(List.append(&2, Backward.Node, newer, node <> rest), 1n+Nat.add(older, List.length(&2, Backward.Node, newer)),
135      older, Backward.at_head(1n+Nat.add(older, List.length(&2, Backward.Node, newer)), older)) ==
136    Sgd.leaf_refusal(node) : Maybe<&2, Sgd.Error>}
137
138# What a parameter's node allows: a node with no operands, a leaf, is
139# accepted; a node with operands is refused with NotLeaf.
140def operand_refusal(targets: List<&2, Nat>) -> Maybe<&2, Sgd.Error>:
141  match targets:
142    case Nil{}:
143      None{}
144    case Con{target, rest}:
145      Some{Sgd.NotLeaf{}}
146
147law sgd_only_leaves_accepted:
148  for node: Backward.Node
149  {Sgd.leaf_refusal(node) == operand_refusal(TensorBackward.targets(node)) : Maybe<&2, Sgd.Error>}
150
151# An id the tape does not hold, the number of its nodes or more, is refused
152# with UnknownNode.
153law sgd_refusal_unknown:
154  for nodes: List<&2, Backward.Node>
155  for +extra: Nat
156  {Sgd.refusal_at(nodes, List.length(&2, Backward.Node, nodes), Nat.add(List.length(&2, Backward.Node, nodes), extra),
157      Backward.at_head(List.length(&2, Backward.Node, nodes), Nat.add(List.length(&2, Backward.Node, nodes), extra))) ==
158    Some{Sgd.UnknownNode{}} : Maybe<&2, Sgd.Error>}
159
160# The learning rate is refused exactly when it is below zero, IEEE `lr < 0`
161# on the bits (`below_zero`, stated by below_zero_words): PyTorch's
162# `lr < 0.0`.
163law sgd_negative_rate_refused:
164  for +lr: F32
165  for below: {Elementwise.below_zero(lr) == True{} : Bool}
166  {Sgd.rate_refusal(lr) == Some{Sgd.InvalidLearningRate{}} : Maybe<&2, Sgd.Error>}
167
168law sgd_rate_accepted:
169  for +lr: F32
170  for above: {Elementwise.below_zero(lr) == False{} : Bool}
171  {Sgd.rate_refusal(lr) == None{} : Maybe<&2, Sgd.Error>}
172
173def from_bits(word: U32) -> F32:
174  U32{bits} = word
175  F32{bits}
176
177def rate_refusals(words: List<&2, U32>) -> List<&2, Maybe<&2, Sgd.Error>>:
178  match words:
179    case Nil{}:
180      Nil{}
181    case Con{word, rest}:
182      Sgd.rate_refusal(from_bits(word)) <> rate_refusals(rest)
183
184# At the edge rates, as PyTorch 2.14.1 answers them: +0.0, -0.0, 0.01, the
185# quiet NaN, a NaN with the sign bit and +inf are accepted; the smallest
186# negative denormal, -0.25, the most negative finite word and -inf are
187# refused.
188law sgd_rate_edges:
189  {rate_refusals([0, 2147483648, 1008981770, 2143289344, 4290772992, 2139095040,
190      2147483649, 3196059648, 4286578687, 4286578688]) ==
191    [None{}, None{}, None{}, None{}, None{}, None{},
192      Some{Sgd.InvalidLearningRate{}}, Some{Sgd.InvalidLearningRate{}}, Some{Sgd.InvalidLearningRate{}},
193      Some{Sgd.InvalidLearningRate{}}] : List<&2, Maybe<&2, Sgd.Error>>}
194
195# A step with a refused parameter leaves the tape unchanged; with every id
196# accepted, a refused learning rate leaves it unchanged too; otherwise the
197# step updates the parameters in list order.
198law sgd_step_refused:
199  for entries: List<&1, Backward.Entry>
200  for +lr: F32
201  for +params: List<&2, Nat>
202  for +error: Sgd.Error
203  for refused: {Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params) ==
204    Some{error} : Maybe<&2, Sgd.Error>}
205  {Sgd.step(Backward.Tape{entries}, lr, params) == (Backward.Tape{entries}, Fail{error}) :
206    Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
207
208law sgd_step_rate_refused:
209  for entries: List<&1, Backward.Entry>
210  for +lr: F32
211  for +params: List<&2, Nat>
212  for +error: Sgd.Error
213  for accepted: {Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params) ==
214    None{} : Maybe<&2, Sgd.Error>}
215  for refused: {Sgd.rate_refusal(lr) == Some{error} : Maybe<&2, Sgd.Error>}
216  {Sgd.step(Backward.Tape{entries}, lr, params) == (Backward.Tape{entries}, Fail{error}) :
217    Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
218
219law sgd_step_accepted:
220  for entries: List<&1, Backward.Entry>
221  for +lr: F32
222  for +params: List<&2, Nat>
223  for accepted: {Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params) ==
224    None{} : Maybe<&2, Sgd.Error>}
225  for rated: {Sgd.rate_refusal(lr) == None{} : Maybe<&2, Sgd.Error>}
226  {Sgd.step(Backward.Tape{entries}, lr, params) ==
227    Sgd.walked_tape(Sgd.descend_all(lr, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, entries)) :
228    Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}