1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise.bend as Elementwise
5	import ./gradient.bend as Gradient
6	import ./tensor_backward.bend as Backward
7	
8	# Plain stochastic gradient descent on the tape, PyTorch's
9	# torch.optim.SGD(params, lr) with no momentum, dampening, nesterov,
10	# weight_decay or maximize. The parameters are a list of tape ids, each a
11	# leaf. step visits them in list order; a parameter whose slot holds g gets
12	# the value p - lr * g, cell by cell on the parallel elementwise build
13	# (Elementwise.Descend), and keeps its slot, so a later read still sees g; a
14	# parameter whose slot is empty, PyTorch's p.grad is None, is left as it was.
15	# A listed id twice is updated twice, as PyTorch's step does. Recording an
16	# operation after a step reads the new values, as a forward after
17	# optimizer.step() reads the updated parameters. Backward.zero_grad empties
18	# the slots.
19	#
20	# step checks every id, then the learning rate, before it updates anything:
21	# an id the tape does not hold is refused with UnknownNode, an id of a node
22	# that is not a leaf with NotLeaf, and a learning rate below zero (IEEE
23	# `lr < 0`, PyTorch's own check: -0.0, NaN and +inf are accepted) with
24	# InvalidLearningRate. A refusal leaves the tape unchanged.
25	
26	type Error is Data:
27	  UnknownNode{}
28	  NotLeaf{}
29	  InvalidLearningRate{}
30	  TensorError{error: Core.Error}
31	
32	def leaf_refusal(node: Backward.Node) -> Maybe<&2, Error>:
33	  match node:
34	    case Backward.Leaf{}:
35	      None{}
36	    case Backward.Add{a, b, left, right}:
37	      Some{NotLeaf{}}
38	    case Backward.Sub{a, b, left, right}:
39	      Some{NotLeaf{}}
40	    case Backward.Mul{a, b, left, right}:
41	      Some{NotLeaf{}}
42	    case Backward.Sum{a, input}:
43	      Some{NotLeaf{}}
44	    case Backward.SumDims{a, input, kept}:
45	      Some{NotLeaf{}}
46	    case Backward.Matmul{a, b, left, right}:
47	      Some{NotLeaf{}}
48	    case Backward.MatmulNT{a, b, left, right}:
49	      Some{NotLeaf{}}
50	    case Backward.Relu{a, input}:
51	      Some{NotLeaf{}}
52	    case Backward.MseLoss{a, b, left, right}:
53	      Some{NotLeaf{}}
54	    case Backward.CrossEntropy{a, input}:
55	      Some{NotLeaf{}}
56	    case Backward.BatchedMatmul{a, b, left, right}:
57	      Some{NotLeaf{}}
58	    case Backward.Slice{a, input, kept, plan}:
59	      Some{NotLeaf{}}
60	    case Backward.Gathered{a, input, plan, entries}:
61	      Some{NotLeaf{}}
62	    case Backward.Constant{}:
63	      None{}
64	    case Backward.Frozen{id, inner}:
65	      leaf_refusal(inner)
66	
67	# Why id cannot be a parameter of a newest-first list of count nodes, or
68	# None when it is a leaf the list holds.
69	def refusal_at(nodes: List<&2, Backward.Node>, +count: Nat, +id: Nat, same: Bool) -> Maybe<&2, Error>:
70	  match nodes count same:
71	    case Nil{} count same:
72	      Some{UnknownNode{}}
73	    case Con{node, rest} 0n same:
74	      Some{UnknownNode{}}
75	    case Con{node, rest} 1n+below True{}:
76	      leaf_refusal(node)
77	    case Con{node, rest} 1n+below False{}:
78	      refusal_at(rest, below, id, Backward.at_head(below, id))
79	
80	def first_refusal(found: Maybe<&2, Error>, later: Maybe<&2, Error>) -> Maybe<&2, Error>:
81	  match found:
82	    case Some{error}:
83	      Some{error}
84	    case None{}:
85	      later
86	
87	# The refusal of the first parameter, in list order, that is not a leaf of
88	# the tape, or None when every one is.
89	def refusal(+nodes: List<&2, Backward.Node>, +count: Nat, params: List<&2, Nat>) -> Maybe<&2, Error>:
90	  match params:
91	    case Nil{}:
92	      None{}
93	    case Con{+id, rest}:
94	      first_refusal(refusal_at(nodes, count, id, Backward.at_head(count, id)), refusal(nodes, count, rest))
95	
96	def Walked() -> Type:
97	  List<&1, Backward.Entry> & Result<&2, &2, Error, Unit>
98	
99	def descended(node: Backward.Node, +shape: Shape.Shape, saved: Backward.Saved,
100	  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Backward.Entry & Result<&2, &2, Error, Unit>:
101	  (operands, result) = pair
102	  (value, g) = operands
103	  match result:
104	    case Fail{error}:
105	      (Backward.Entry{node, shape, value, saved, Gradient.Held{g}}, Fail{TensorError{error}})
106	    case Done{updated}:
107	      (Backward.Entry{node, shape, updated, saved, Gradient.Held{g}}, Done{Unit{}})
108	
109	# One parameter: p - lr * g, cell by cell, when its slot holds g; unchanged
110	# when the slot is empty.
111	def descend(+lr: F32, entry: Backward.Entry) -> Backward.Entry & Result<&2, &2, Error, Unit>:
112	  Backward.Entry{node, +shape, value, saved, slot} = entry
113	  match slot:
114	    case Gradient.Unset{}:
115	      (Backward.Entry{node, shape, value, saved, Gradient.Unset{}}, Done{Unit{}})
116	    case Gradient.Held{g}:
117	      descended(node, shape, saved, Elementwise.binary(Elementwise.Descend{lr}, value, g))
118	
119	def placed(rest: List<&1, Backward.Entry>, pair: Backward.Entry & Result<&2, &2, Error, Unit>) -> Walked():
120	  (entry, result) = pair
121	  (entry <> rest, result)
122	
123	def kept(entry: Backward.Entry, pair: Walked()) -> Walked():
124	  (rest, result) = pair
125	  (entry <> rest, result)
126	
127	# Updates entry id of a newest-first list of count entries and keeps every
128	# other entry.
129	def descend_at(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +id: Nat, same: Bool) -> Walked():
130	  match entries count same:
131	    case Nil{} count same:
132	      (Nil{}, Fail{UnknownNode{}})
133	    case Con{entry, rest} 0n same:
134	      (entry <> rest, Fail{UnknownNode{}})
135	    case Con{entry, rest} 1n+below True{}:
136	      placed(rest, descend(lr, entry))
137	    case Con{entry, rest} 1n+below False{}:
138	      kept(entry, descend_at(lr, rest, below, id, Backward.at_head(below, id)))
139	
140	def continued(rec: List<&1, Backward.Entry> -> Walked(), pair: Walked()) -> Walked():
141	  (entries, result) = pair
142	  match result:
143	    case Fail{error}:
144	      (entries, Fail{error})
145	    case Done{unit}:
146	      rec(entries)
147	
148	# Each parameter in list order, each on the entries the one before it left.
149	def descend_all(+lr: F32, +count: Nat, params: List<&2, Nat>, entries: List<&1, Backward.Entry>) -> Walked():
150	  match params:
151	    case Nil{}:
152	      (entries, Done{Unit{}})
153	    case Con{+id, rest}:
154	      continued(entries => descend_all(lr, count, rest, entries), descend_at(lr, entries, count, id, Backward.at_head(count, id)))
155	
156	def walked_tape(pair: Walked()) -> Backward.Tape & Result<&2, &2, Error, Unit>:
157	  (entries, result) = pair
158	  (Backward.Tape{entries}, result)
159	
160	def refused_if(below: Bool) -> Maybe<&2, Error>:
161	  match below:
162	    case True{}:
163	      Some{InvalidLearningRate{}}
164	    case False{}:
165	      None{}
166	
167	# InvalidLearningRate when lr is below zero, as PyTorch's `lr < 0.0`; None
168	# for every other word, -0.0, NaN and +inf included.
169	def rate_refusal(+lr: F32) -> Maybe<&2, Error>:
170	  refused_if(Elementwise.below_zero(lr))
171	
172	def rated(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +params: List<&2, Nat>, refused: Maybe<&2, Error>) ->
173	  Backward.Tape & Result<&2, &2, Error, Unit>:
174	  match refused:
175	    case Some{error}:
176	      (Backward.Tape{entries}, Fail{error})
177	    case None{}:
178	      walked_tape(descend_all(lr, count, params, entries))
179	
180	def checked(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +params: List<&2, Nat>, refused: Maybe<&2, Error>) ->
181	  Backward.Tape & Result<&2, &2, Error, Unit>:
182	  match refused:
183	    case Some{error}:
184	      (Backward.Tape{entries}, Fail{error})
185	    case None{}:
186	      rated(lr, entries, count, params, rate_refusal(lr))
187	
188	def planned(+lr: F32, +params: List<&2, Nat>, pair: List<&1, Backward.Entry> & List<&2, Backward.Node>) ->
189	  Backward.Tape & Result<&2, &2, Error, Unit>:
190	  (entries, +nodes) = pair
191	  +count = List.length(&2, Backward.Node, nodes)
192	  checked(lr, entries, count, params, refusal(nodes, count, params))
193	
194	# One SGD step with learning rate lr over the parameters params.
195	def step(tape: Backward.Tape, +lr: F32, +params: List<&2, Nat>) -> Backward.Tape & Result<&2, &2, Error, Unit>:
196	  Backward.Tape{entries} = tape
197	  planned(lr, params, Backward.listed(entries))
198	
199	def error_code(error: Error) -> U32:
200	  match error:
201	    case UnknownNode{}:
202	      Backward.error_code(Backward.UnknownNode{})
203	    case NotLeaf{}:
204	      109
205	    case InvalidLearningRate{}:
206	      110
207	    case TensorError{error}:
208	      Core.error_code(error)