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/elementwise.bend as Elementwise
6	import ../tensor/lifecycle.bend as Lifecycle
7	import ../tensor/contracts/replayed.bend as Replayed
8	import ../tensor/contracts/LAWS.bend as Contracts
9	import ../tensor/contracts/binary_LAWS.bend as Binary
10	import ./gradient.bend as Gradient
11	import ./tensor_backward.bend as Backward
12	import ./tensor_backward_LAWS.bend as TensorBackward
13	import ./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.
17	law 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.
25	law 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.
39	def 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.
68	law 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.
88	def 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.
97	law 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.
110	law 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	
116	law 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.
128	law 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.
140	def 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	
147	law 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.
153	law 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`.
163	law 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	
168	law 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	
173	def from_bits(word: U32) -> F32:
174	  U32{bits} = word
175	  F32{bits}
176	
177	def 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.
188	law 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.
198	law 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	
208	law 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	
219	law 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>}