src/tensor/loss_ops_LAWS.bend

bendtorch-core@0.1.0.0 · 98 lines · 12 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
3
4# The elementwise operations of the mean and of the two losses, each pinned
5# to the F32 operations it applies, in PyTorch's order. F32.div, F32.exp and
6# F32.log are opaque in Base, so these laws state which operations run on
7# which words, not how close exp and log are to the exact values.
8
9# mean: the sum divided by the count.
10law scalar_divide:
11  for +left: F32
12  for +right: F32
13  {Elementwise.scalar(Elementwise.Divide{}, left, right) == F32.div(left, right) : F32}
14
15# exp of the left word, the right word unread.
16law scalar_exp_left:
17  for +left: F32
18  for +right: F32
19  {Elementwise.scalar(Elementwise.ExpLeft{}, left, right) == F32.exp(left) : F32}
20
21# log of the left word, the right word unread.
22law scalar_log_left:
23  for +left: F32
24  for +right: F32
25  {Elementwise.scalar(Elementwise.LogLeft{}, left, right) == F32.log(left) : F32}
26
27# One-hot: 1.0 where the two words are the same bits, +0.0 elsewhere.
28law scalar_matches_same:
29  for +left: F32
30  for +right: F32
31  for same: {U32.is_eq(Elementwise.bits(left), Elementwise.bits(right)) == True{} : Bool}
32  {Elementwise.scalar(Elementwise.Matches{}, left, right) == 1.0 : F32}
33
34law scalar_matches_other:
35  for +left: F32
36  for +right: F32
37  for other: {U32.is_eq(Elementwise.bits(left), Elementwise.bits(right)) == False{} : Bool}
38  {Elementwise.scalar(Elementwise.Matches{}, left, right) == 0.0 : F32}
39
40# The nll term: -lp where the one-hot word is 1.0 (0x3F800000), +0.0
41# elsewhere.
42law scalar_pick_negated_target:
43  for +left: F32
44  for +right: F32
45  for target: {Elementwise.is_one(right) == True{} : Bool}
46  {Elementwise.scalar(Elementwise.PickNegated{}, left, right) == F32.neg(left) : F32}
47
48law scalar_pick_negated_other:
49  for +left: F32
50  for +right: F32
51  for other: {Elementwise.is_one(right) == False{} : Bool}
52  {Elementwise.scalar(Elementwise.PickNegated{}, left, right) == 0.0 : F32}
53
54# The three backward operations take the gradient g as the left word, as
55# `Mul` does in the Mul VJP, and a saved word as the right one.
56
57# PyTorch's mse_loss backward: (scale * d) * g, scale = 2 / N, for g on the
58# left and d on the right.
59law scalar_mse_grad:
60  for +scale: F32
61  for +left: F32
62  for +right: F32
63  {Elementwise.scalar(Elementwise.MseGrad{scale}, left, right) == F32.mul(F32.mul(scale, right), left) : F32}
64
65# PyTorch's nll_loss backward with mean reduction, for g on the left and a
66# one-hot word on the right: -(g / N) at a target, +0.0 elsewhere.
67law scalar_one_hot_grad_target:
68  for +count: F32
69  for +left: F32
70  for +right: F32
71  for target: {Elementwise.is_one(right) == True{} : Bool}
72  {Elementwise.scalar(Elementwise.OneHotGrad{count}, left, right) == F32.neg(F32.div(left, count)) : F32}
73
74law scalar_one_hot_grad_other:
75  for +count: F32
76  for +left: F32
77  for +right: F32
78  for other: {Elementwise.is_one(right) == False{} : Bool}
79  {Elementwise.scalar(Elementwise.OneHotGrad{count}, left, right) == 0.0 : F32}
80
81# The rest of PyTorch's log_softmax backward, for g on the left and lp on
82# the right: -(exp(lp) * r) with r the row sum of the one-hot gradient. In a
83# row of one column (`single`) the sum is the target word itself,
84# r = -(g / N), -0.0 kept as PyTorch keeps it.
85law scalar_softmax_grad_one_class:
86  for +count: F32
87  for +left: F32
88  for +right: F32
89  {Elementwise.scalar(Elementwise.SoftmaxGrad{count, True{}}, left, right) ==
90    F32.neg(F32.mul(F32.exp(right), F32.neg(F32.div(left, count)))) : F32}
91
92# In a wider row the +0.0 words are added: r = -(g / N) + 0.0.
93law scalar_softmax_grad_classes:
94  for +count: F32
95  for +left: F32
96  for +right: F32
97  {Elementwise.scalar(Elementwise.SoftmaxGrad{count, False{}}, left, right) ==
98    F32.neg(F32.mul(F32.exp(right), F32.add(F32.neg(F32.div(left, count)), 0.0))) : F32}