1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/main.bend as Storage
4	import ../storage/model.bend as Model
5	import ../tensor/core.bend as Core
6	import ../tensor/elementwise.bend as Elementwise
7	import ../tensor/matmul.bend as Matrix
8	import ../tensor/parallel.bend as Parallel
9	import ../tensor/reduction.bend as Reduction
10	import ./reduce.bend as Reduce
11	import ./sum_to_size.bend as SumToSize
12	
13	# Loss functions with mean reduction, as PyTorch's mse_loss and
14	# cross_entropy over class logits. Every step is a tensor operation that
15	# builds its output over cells in parallel: Elementwise.binary,
16	# sum_to_size and the row maximum below. The mean is the ordered sum of
17	# Reduction.sum_dynamic, added left to right from +0.0, divided by the
18	# count with F32.div.
19	
20	type Error is Data:
21	  TensorError{error: Core.Error}
22	  ReduceError{error: Reduce.Error}
23	  NotLogits{}
24	  TargetOutOfRange{}
25	
26	def error_code(error: Error) -> U32:
27	  match error:
28	    case TensorError{error}:
29	      Core.error_code(error)
30	    case ReduceError{error}:
31	      Reduce.error_code(error)
32	    case NotLogits{}:
33	      107
34	    case TargetOutOfRange{}:
35	      108
36	
37	# A loss next to the two tensors its backward reads.
38	type Recorded is Type:
39	  Recorded{loss: Core.Dynamic, left: Core.Dynamic, right: Core.Dynamic}
40	
41	# A count as an F32, rounded as PyTorch converts an integer count.
42	def count_f32(+count: Nat) -> F32:
43	  U32.to_f32(U32.from_nat(count))
44	
45	# 2 / N, the factor of PyTorch's mse_loss backward with mean reduction, for
46	# the N cells of the broadcast shape.
47	def mse_scale(+shape: Shape.Shape) -> F32:
48	  F32.div(2.0, count_f32(Shape.logical_numel(Shape.dimensions(shape))))
49	
50	# The first dimension of [rows, columns] logits, the batch.
51	def first_dimension(+dimensions: List<&2, U32>) -> Nat:
52	  match dimensions:
53	    case Nil{}:
54	      0n
55	    case Con{+first, +rest}:
56	      U32.to_nat(first)
57	
58	def rows(+shape: Shape.Shape) -> Nat:
59	  first_dimension(Shape.dimensions(shape))
60	
61	# Whether [rows, columns] logits have one column, one class.
62	def single_column(+dimensions: List<&2, U32>) -> Bool:
63	  match dimensions:
64	    case [+rows, +columns]:
65	      U32.is_eq(columns, 1)
66	    case _:
67	      False{}
68	
69	def one_class(+shape: Shape.Shape) -> Bool:
70	  single_column(Shape.dimensions(shape))
71	
72	# Row maximum.
73	
74	# PyTorch's maximum: NaN when either word is NaN, else the larger word, and
75	# on a tie the later word, as `_mm512_max_ps(m, v)` keeps its second operand.
76	def max_step(+m: F32, +v: F32) -> F32:
77	  Bool.pick(F32, Bool.or(F32.is_gt(m, v), F32.is_ne(m, m)), m, v)
78	
79	def row_next(+m: F32, rec: F32 -> Array<F32> -> F32 & Array<F32>, pair: Array<F32> & F32) -> F32 & Array<F32>:
80	  (array, v) = pair
81	  rec(max_step(m, v), array)
82	
83	# max_step of m and the `fuel` words from `position` on, left to right.
84	def row_rest(fuel: Nat, +position: Nat, +m: F32, array: Array<F32>) -> F32 & Array<F32>:
85	  match fuel:
86	    case 0n:
87	      (m, array)
88	    case 1n+remaining:
89	      row_next(m, value => kept => row_rest(remaining, 1n+position, value, kept), Array.get(F32, array, U32.from_nat(position)))
90	
91	def row_first(+columns: Nat, +start: Nat, pair: Array<F32> & F32) -> F32 & Array<F32>:
92	  (array, first) = pair
93	  row_rest(Nat.sub(columns, 1n), 1n+start, first, array)
94	
95	# The F32 of a word.
96	def from_word(word: U32) -> F32:
97	  U32{bits} = word
98	  F32{bits}
99	
100	# The maximum of row `row` of a row-major [rows, columns] array: its first
101	# word, then max_step with each later word, left to right. A row of no
102	# words is -inf (0xFF800000).
103	def row_max(+columns: Nat, +row: Nat, array: Array<F32>) -> F32 & Array<F32>:
104	  match columns:
105	    case 0n:
106	      (from_word(4286578688), array)
107	    case 1n+less:
108	      row_first(1n+less, Nat.mul(row, 1n+less), Array.get(F32, array, U32.from_nat(Nat.mul(row, 1n+less))))
109	
110	type Width is Data:
111	  Width{columns: Nat}
112	
113	def width(+given: Width) -> Width:
114	  Width{+columns} = given
115	  Width{columns}
116	
117	# The cell function of the row-maximum build: cell `position` is the
118	# maximum of row `position`. The state is the input array next to a Unit, as
119	# in Reduce.build.
120	def max_cell(+width: Width, +position: Nat, state: Array<F32> & Unit, inside: Bool) -> F32 & (Array<F32> & Unit):
121	  Width{+columns} = width
122	  (array, none) = state
123	  match inside:
124	    case False{}:
125	      (0.0, (array, none))
126	    case True{}:
127	      Reduce.keep(row_max(columns, position, array), none)
128	
129	def max_build(+depth: Nat, +splits: Nat, +columns: Nat, +count: Nat, +base: Nat, array: Array<F32>) -> Array<F32>:
130	  Parallel.build(~Width, ~Width, ~Array<F32>, ~Unit, ~width, ~Reduce.split, ~max_cell, Width{columns}, depth, splits, count, base,
131	    (array, Unit{}))
132	
133	def max_balanced(+depth: Nat, +splits: Nat, +columns: Nat, +count: Nat, +base: Nat, array: Array<F32>) ->
134	  Model.Balanced(F32, depth, max_build(depth, splits, columns, count, base, array)):
135	  Parallel.balanced(~Width, ~Width, ~Array<F32>, ~Unit, ~width, ~Reduce.split, ~max_cell, Width{columns}, depth, splits, count, base,
136	    (array, Unit{}))
137	
138	def max_rows(+target: Shape.Shape, +columns: Nat, +depth: Nat, +length: Nat,
139	  -valid: {Storage.bounds(depth, length) == True{} : Bool},
140	  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
141	  taken: Core.Dynamic & Array<F32>) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
142	  (kept, array) = taken
143	  +splits = Matrix.splits(depth, columns)
144	  (kept, Done{Core.Packed{target, Core.TTensor{Storage.Owned{depth, length,
145	    max_build(depth, splits, columns, length, 0n, array),
146	    max_balanced(depth, splits, columns, length, 0n, array),
147	    valid}, exact}}})
148	
149	def max_allocated(input: Core.Dynamic, +target: Shape.Shape, +columns: Nat,
150	  result: Result<&2, &1, Core.Error, Core.Tensor<target>>) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
151	  match result:
152	    case Fail{error}:
153	      (input, Fail{error})
154	    case Done{Core.TTensor{storage, exact}}:
155	      match storage:
156	        case Storage.Empty{}:
157	          (input, Done{Core.Packed{target, Core.TTensor{Storage.Empty{}, exact}}})
158	        case Storage.Owned{+depth, +length, array, balanced, valid}:
159	          max_rows(target, columns, depth, length, valid, exact, Elementwise.take(input))
160	
161	# The maximum of each row of [rows, columns] logits, a [rows, 1] tensor.
162	# Each output cell walks its own row, so the cells are built in parallel,
163	# forked by the work of a row as Matrix.splits forks a matmul. The input
164	# comes back unchanged.
165	def row_maxes(input: Core.Dynamic, +target: Shape.Shape, +columns: Nat) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
166	  max_allocated(input, target, columns, Core.new(target, 0.0))
167	
168	# Mean.
169	
170	# The result of a binary operation whose operands the caller no longer needs.
171	def quotient(pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Result<&2, &1, Core.Error, Core.Dynamic>:
172	  (operands, result) = pair
173	  result
174	
175	def divided_by(+count: Nat, pair: Core.Dynamic & Core.Dynamic,
176	  divisor: Result<&2, &1, Core.Error, Core.Tensor<Reduction.scalar_shape()>>) -> Result<&2, &1, Core.Error, Core.Dynamic>:
177	  (input, total) = pair
178	  match divisor:
179	    case Fail{error}:
180	      Fail{error}
181	    case Done{tensor}:
182	      quotient(Elementwise.binary(Elementwise.Divide{}, total, Core.pack(Reduction.scalar_shape(), tensor)))
183	
184	# The sum of every word, in ascending flat order from +0.0
185	# (Reduction.sum_dynamic), divided by `count`: a rank-0 tensor.
186	def total_over(+count: Nat, result: Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>) -> Result<&2, &1, Core.Error, Core.Dynamic>:
187	  match result:
188	    case Fail{error}:
189	      Fail{error}
190	    case Done{pair}:
191	      divided_by(count, pair, Core.new(Reduction.scalar_shape(), count_f32(count)))
192	
193	def mean_shaped(pair: Core.Dynamic & Shape.Shape) -> Result<&2, &1, Core.Error, Core.Dynamic>:
194	  (input, +shape) = pair
195	  total_over(Shape.logical_numel(Shape.dimensions(shape)), Reduction.sum_dynamic(input))
196	
197	# PyTorch's mean of every element: the ordered sum divided by the element
198	# count, F32.div of the two. The sum adds left to right from +0.0, where
199	# PyTorch's CPU sum adds in a cascade, so for many words the last bits can
200	# differ.
201	def mean(input: Core.Dynamic) -> Result<&2, &1, Core.Error, Core.Dynamic>:
202	  mean_shaped(Core.shape(input))
203	
204	# mse_loss.
205	
206	# The loss next to the two tensors the backward reads.
207	def recorded(left: Core.Dynamic, right: Core.Dynamic, result: Result<&2, &1, Core.Error, Core.Dynamic>) -> Result<&2, &1, Error, Recorded>:
208	  match result:
209	    case Fail{error}:
210	      Fail{TensorError{error}}
211	    case Done{loss}:
212	      Done{Recorded{loss, left, right}}
213	
214	def mse_squared(e: Core.Dynamic, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
215	  Result<&2, &1, Error, Recorded>:
216	  (operands, result) = pair
217	  (kept, copy) = operands
218	  match result:
219	    case Fail{error}:
220	      Fail{TensorError{error}}
221	    case Done{squares}:
222	      recorded(e, kept, mean(squares))
223	
224	def mse_copied(e: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) -> Result<&2, &1, Error, Recorded>:
225	  (d, copy) = pair
226	  mse_squared(e, Elementwise.binary(Elementwise.Mul{}, d, copy))
227	
228	def mse_both(input: Core.Dynamic, target: Core.Dynamic, d: Result<&2, &1, Core.Error, Core.Dynamic>,
229	  e: Result<&2, &1, Core.Error, Core.Dynamic>) -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Error, Recorded>:
230	  match d e:
231	    case Fail{error} e:
232	      ((input, target), Fail{TensorError{error}})
233	    case Done{d} Fail{error}:
234	      ((input, target), Fail{TensorError{error}})
235	    case Done{d} Done{e}:
236	      ((input, target), mse_copied(e, Core.clone_dynamic(d)))
237	
238	def mse_differences(d: Result<&2, &1, Core.Error, Core.Dynamic>, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
239	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Error, Recorded>:
240	  (operands, e) = pair
241	  (target, input) = operands
242	  mse_both(input, target, d, e)
243	
244	def mse_first(pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
245	  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Error, Recorded>:
246	  (operands, d) = pair
247	  (input, target) = operands
248	  mse_differences(d, Elementwise.binary(Elementwise.Sub{}, target, input))
249	
250	# PyTorch's mse_loss(input, target) with mean reduction: d = input - target
251	# and e = target - input over the broadcast shape, then the mean of d * d.
252	# Both inputs come back unchanged, next to the loss and e and d (in that
253	# order, the order the tape saves them), the two differences the backward
254	# scales: PyTorch's backward computes
255	# (2 / N * d) * g for the input and (2 / N * e) * g for the target.
256	def mse_loss(input: Core.Dynamic, target: Core.Dynamic) -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Error, Recorded>:
257	  mse_first(Elementwise.binary(Elementwise.Sub{}, input, target))
258	
259	# cross_entropy.
260	
261	# Whether every target is below `columns`, and there are `rows` of them.
262	def targets_fit(+targets: List<&2, U32>, +rows: Nat, +columns: Nat) -> Bool:
263	  match targets rows:
264	    case Nil{} 0n:
265	      True{}
266	    case Con{+target, +rest} 1n+below:
267	      Bool.and(Nat.is_lt(U32.to_nat(target), columns), targets_fit(rest, below, columns))
268	    case _ _:
269	      False{}
270	
271	def as_words(+targets: List<&2, U32>) -> List<&2, F32>:
272	  match targets:
273	    case Nil{}:
274	      Nil{}
275	    case Con{+target, +rest}:
276	      U32.to_f32(target) <> as_words(rest)
277	
278	# The class numbers 0.0, 1.0, ... of `fuel` columns from `next`.
279	def classes(fuel: Nat, +next: Nat) -> List<&2, F32>:
280	  match fuel:
281	    case 0n:
282	      Nil{}
283	    case 1n+remaining:
284	      count_f32(next) <> classes(remaining, 1n+next)
285	
286	def one_hot_pair(targets: Result<&2, &1, Core.Error, Core.Dynamic>, columns: Result<&2, &1, Core.Error, Core.Dynamic>) ->
287	  Result<&2, &1, Core.Error, Core.Dynamic>:
288	  match targets columns:
289	    case Fail{error} columns:
290	      Fail{error}
291	    case Done{targets} Fail{error}:
292	      Fail{error}
293	    case Done{targets} Done{columns}:
294	      quotient(Elementwise.binary(Elementwise.Matches{}, targets, columns))
295	
296	# The one-hot rows of the targets, [rows, columns]: 1.0 where the column is
297	# the row's target, +0.0 elsewhere. The [rows, 1] targets and the [1, columns]
298	# class numbers broadcast against each other with Matches, which compares
299	# their words. With no rows the one-hot rows are empty, and the class
300	# numbers are not built: [1, columns] could exceed the capacity that
301	# [0, columns] does not.
302	def one_hot(+targets: List<&2, U32>, +rows: Nat, +columns: Nat) -> Result<&2, &1, Core.Error, Core.Dynamic>:
303	  match rows:
304	    case 0n:
305	      Core.construct([0, U32.from_nat(columns)], Nil{})
306	    case 1n+less:
307	      one_hot_pair(Core.construct([U32.from_nat(1n+less), 1], as_words(targets)),
308	        Core.construct([1, U32.from_nat(columns)], classes(columns, 0n)))
309	
310	def ce_total(+rows: Nat, lp: Core.Dynamic, onehot: Core.Dynamic, result: Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>) ->
311	  Result<&2, &1, Error, Recorded>:
312	  recorded(lp, onehot, total_over(rows, result))
313	
314	def ce_picked(+rows: Nat, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Result<&2, &1, Error, Recorded>:
315	  (operands, result) = pair
316	  (lp, onehot) = operands
317	  match result:
318	    case Fail{error}:
319	      Fail{TensorError{error}}
320	    case Done{picked}:
321	      ce_total(rows, lp, onehot, Reduction.sum_dynamic(picked))
322	
323	def ce_logged(+rows: Nat, onehot: Core.Dynamic, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
324	  Result<&2, &1, Error, Recorded>:
325	  (operands, result) = pair
326	  (shifted, lse) = operands
327	  match result:
328	    case Fail{error}:
329	      Fail{TensorError{error}}
330	    case Done{lp}:
331	      ce_picked(rows, Elementwise.binary(Elementwise.PickNegated{}, lp, onehot))
332	
333	def ce_lse(+rows: Nat, onehot: Core.Dynamic, shifted: Core.Dynamic, result: Result<&2, &1, Core.Error, Core.Dynamic>) ->
334	  Result<&2, &1, Error, Recorded>:
335	  match result:
336	    case Fail{error}:
337	      Fail{TensorError{error}}
338	    case Done{lse}:
339	      ce_logged(rows, onehot, Elementwise.binary(Elementwise.Sub{}, shifted, lse))
340	
341	def ce_summed(+rows: Nat, onehot: Core.Dynamic, shifted: Core.Dynamic, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>,
342	  zero: Result<&2, &1, Core.Error, Core.Tensor<Reduction.scalar_shape()>>) -> Result<&2, &1, Error, Recorded>:
343	  (exps, result) = pair
344	  match result zero:
345	    case Fail{error} zero:
346	      Fail{ReduceError{error}}
347	    case Done{sums} Fail{error}:
348	      Fail{TensorError{error}}
349	    case Done{sums} Done{tensor}:
350	      ce_lse(rows, onehot, shifted,
351	        quotient(Elementwise.binary(Elementwise.LogLeft{}, sums, Core.pack(Reduction.scalar_shape(), tensor))))
352	
353	def ce_exps(+rows: Nat, onehot: Core.Dynamic, +target: Shape.Shape, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
354	  Result<&2, &1, Error, Recorded>:
355	  (operands, result) = pair
356	  (shifted, zero) = operands
357	  match result:
358	    case Fail{error}:
359	      Fail{TensorError{error}}
360	    case Done{exps}:
361	      ce_summed(rows, onehot, shifted, SumToSize.sum_to_size(exps, target), Core.new(Reduction.scalar_shape(), 0.0))
362	
363	def ce_shifted(+rows: Nat, onehot: Core.Dynamic, +target: Shape.Shape, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
364	  zero: Result<&2, &1, Core.Error, Core.Tensor<Reduction.scalar_shape()>>) -> Core.Dynamic & Result<&2, &1, Error, Recorded>:
365	  (operands, result) = pair
366	  (logits, maxes) = operands
367	  match result zero:
368	    case Fail{error} zero:
369	      (logits, Fail{TensorError{error}})
370	    case Done{shifted} Fail{error}:
371	      (logits, Fail{TensorError{error}})
372	    case Done{shifted} Done{tensor}:
373	      (logits, ce_exps(rows, onehot, target,
374	        Elementwise.binary(Elementwise.ExpLeft{}, shifted, Core.pack(Reduction.scalar_shape(), tensor))))
375	
376	def ce_maxed(+rows: Nat, onehot: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>) ->
377	  Core.Dynamic & Result<&2, &1, Error, Recorded>:
378	  (logits, result) = pair
379	  match result:
380	    case Fail{error}:
381	      (logits, Fail{TensorError{error}})
382	    case Done{maxes}:
383	      ce_shifted(rows, onehot, target, Elementwise.binary(Elementwise.Sub{}, logits, maxes), Core.new(Reduction.scalar_shape(), 0.0))
384	
385	def ce_targeted(logits: Core.Dynamic, +rows: Nat, +columns: Nat, onehot: Result<&2, &1, Core.Error, Core.Dynamic>,
386	  result: Shape.ConstructResult<[U32.from_nat(rows), 1]>) -> Core.Dynamic & Result<&2, &1, Error, Recorded>:
387	  match onehot result:
388	    case Fail{error} result:
389	      (logits, Fail{TensorError{error}})
390	    case Done{onehot} Shape.SError{error}:
391	      (logits, Fail{TensorError{Core.ShapeError{error}}})
392	    case Done{onehot} Shape.SBuilt{+target, preserved}:
393	      ce_maxed(rows, onehot, target, row_maxes(logits, target, columns))
394	
395	def ce_checked(logits: Core.Dynamic, +targets: List<&2, U32>, +rows: Nat, +columns: Nat, fits: Bool) ->
396	  Core.Dynamic & Result<&2, &1, Error, Recorded>:
397	  match fits:
398	    case False{}:
399	      (logits, Fail{TargetOutOfRange{}})
400	    case True{}:
401	      ce_targeted(logits, rows, columns, one_hot(targets, rows, columns), Shape.construct([U32.from_nat(rows), 1]))
402	
403	def ce_dimensioned(logits: Core.Dynamic, +targets: List<&2, U32>, +dimensions: List<&2, U32>) ->
404	  Core.Dynamic & Result<&2, &1, Error, Recorded>:
405	  match dimensions:
406	    case [+rows, +columns]:
407	      ce_checked(logits, targets, U32.to_nat(rows), U32.to_nat(columns), targets_fit(targets, U32.to_nat(rows), U32.to_nat(columns)))
408	    case _:
409	      (logits, Fail{NotLogits{}})
410	
411	def ce_shaped(+targets: List<&2, U32>, pair: Core.Dynamic & Shape.Shape) -> Core.Dynamic & Result<&2, &1, Error, Recorded>:
412	  (logits, +shape) = pair
413	  ce_dimensioned(logits, targets, Shape.dimensions(shape))
414	
415	# PyTorch's cross_entropy(logits, targets) with mean reduction, for logits
416	# of [rows, columns] and one class index below `columns` per row. Per row,
417	# the stable log_softmax: m = the row maximum (max_step), shifted = x - m,
418	# lse = log of the sum of exp(shifted) in ascending order from +0.0, and
419	# lp = shifted - lse, PyTorch's (x - max) - log(sum). The loss is the sum of
420	# -lp at each row's target, in row order from +0.0, divided by the rows.
421	# The logits come back unchanged, next to the loss, lp and the one-hot rows
422	# of the targets, which the backward reads. Logits that are not of rank 2
423	# fail with NotLogits (107); a target count other than the rows or a target
424	# at or past `columns` with TargetOutOfRange (108).
425	def cross_entropy(logits: Core.Dynamic, +targets: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Recorded>:
426	  ce_shaped(targets, Core.shape(logits))