src/autograd/loss.bend

bendtorch-core@0.1.0.0 · 426 lines · 53 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/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
7import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
8import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
9import bendtorch-core@0.1.0.0/src/tensor/reduction.bend as Reduction
10import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
11import bendtorch-core@0.1.0.0/src/autograd/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
20type Error is Data:
21  TensorError{error: Core.Error}
22  ReduceError{error: Reduce.Error}
23  NotLogits{}
24  TargetOutOfRange{}
25
26def 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.
38type 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.
42def 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.
47def 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.
51def 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
58def rows(+shape: Shape.Shape) -> Nat:
59  first_dimension(Shape.dimensions(shape))
60
61# Whether [rows, columns] logits have one column, one class.
62def 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
69def 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.
76def 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
79def 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.
84def 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
91def 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.
96def 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).
103def 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
110type Width is Data:
111  Width{columns: Nat}
112
113def 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.
120def 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
129def 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
133def 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
138def 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
149def 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.
165def 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.
171def 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
175def 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.
186def 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
193def 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.
201def 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.
207def 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
214def 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
224def 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
228def 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
238def 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
244def 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.
256def 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.
262def 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
271def 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`.
279def 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
286def 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.
302def 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
310def 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
314def 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
323def 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
333def 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
341def 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
353def 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
363def 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
376def 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
385def 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
395def 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
403def 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
411def 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).
425def cross_entropy(logits: Core.Dynamic, +targets: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Recorded>:
426  ce_shaped(targets, Core.shape(logits))