src/tensor/relu_LAWS.bend

bendtorch-core@0.1.0.0 · 117 lines · 15 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
3
4# ReLU and its gradient as elementwise operations, and the two F32
5# comparisons with zero they need, pinned on the IEEE 754 bits.
6
7def from_bits(word: U32) -> F32:
8  U32{bits} = word
9  F32{bits}
10
11# The edge words: +0.0, -0.0, the smallest and largest denormals, the
12# smallest normal, 1.0, the largest finite, the infinities, the signaling
13# NaN 0x7F800001, the quiet NaN 0x7FC00000 and the all-ones NaN 0x7FFFFFFF,
14# each followed by its negative.
15def edges() -> List<&2, U32>:
16  [0, 2147483648, 1, 2147483649, 8388607, 2155872255, 8388608, 2155872256, 1065353216, 3212836864,
17    2139095039, 4286578687, 2139095040, 4286578688, 2139095041, 4286578689, 2143289344, 4290772992,
18    2147483647, 4294967295]
19
20def below_zero_at(words: List<&2, U32>) -> List<&2, Bool>:
21  match words:
22    case Nil{}:
23      Nil{}
24    case Con{word, rest}:
25      Elementwise.below_zero(from_bits(word)) <> below_zero_at(rest)
26
27def at_most_zero_at(words: List<&2, U32>) -> List<&2, Bool>:
28  match words:
29    case Nil{}:
30      Nil{}
31    case Con{word, rest}:
32      Elementwise.at_most_zero(from_bits(word)) <> at_most_zero_at(rest)
33
34# The output words of `operation` with each word as the left operand and
35# `right` as the right one.
36def scalar_at(+operation: Elementwise.Operation, +right: U32, words: List<&2, U32>) -> List<&2, U32>:
37  match words:
38    case Nil{}:
39      Nil{}
40    case Con{word, rest}:
41      Elementwise.bits(Elementwise.scalar(operation, from_bits(word), from_bits(right))) <> scalar_at(operation, right, rest)
42
43# `below_zero` is IEEE 754 `x < 0`: true exactly on the words above -0.0
44# (0x80000000) up to -inf (0xFF800000). +0.0, -0.0 and every NaN, of either
45# sign, are not below zero.
46law below_zero_words:
47  for +x: F32
48  {Elementwise.below_zero(x) ==
49    Bool.and(U32.is_lt(2147483648, Elementwise.bits(x)), U32.is_le(Elementwise.bits(x), 4286578688)) : Bool}
50
51# `at_most_zero` is IEEE 754 `x <= 0`: true exactly on +0.0 (0x00000000) and
52# the words from -0.0 (0x80000000) up to -inf (0xFF800000). No NaN is at most
53# zero.
54law at_most_zero_words:
55  for +x: F32
56  {Elementwise.at_most_zero(x) ==
57    Bool.or(U32.is_eq(Elementwise.bits(x), 0),
58      Bool.and(U32.is_le(2147483648, Elementwise.bits(x)), U32.is_le(Elementwise.bits(x), 4286578688))) : Bool}
59
60# The two comparisons at the edge words, as PyTorch's `x < 0` and `x <= 0`
61# answer them.
62law below_zero_edges:
63  {below_zero_at(edges()) ==
64    [False{}, False{}, False{}, True{}, False{}, True{}, False{}, True{}, False{}, True{},
65      False{}, True{}, False{}, True{}, False{}, False{}, False{}, False{}, False{}, False{}] : List<&2, Bool>}
66
67law at_most_zero_edges:
68  {at_most_zero_at(edges()) ==
69    [True{}, True{}, False{}, True{}, False{}, True{}, False{}, True{}, False{}, True{},
70      False{}, True{}, False{}, True{}, False{}, False{}, False{}, False{}, False{}, False{}] : List<&2, Bool>}
71
72# ReluLeft is PyTorch's relu, clamp_min(x, 0), on the left word: +0.0 where
73# the word is below zero, and the word itself otherwise, so -0.0 stays -0.0
74# and a NaN keeps its sign and payload. The right word is ignored: the
75# forward is the pointwise product with a rank-0 +0.0.
76law scalar_relu_left_below:
77  for +left: F32
78  for +right: F32
79  for below: {Elementwise.below_zero(left) == True{} : Bool}
80  {Elementwise.scalar(Elementwise.ReluLeft{}, left, right) == 0.0 : F32}
81
82law scalar_relu_left_kept:
83  for +left: F32
84  for +right: F32
85  for kept: {Elementwise.below_zero(left) == False{} : Bool}
86  {Elementwise.scalar(Elementwise.ReluLeft{}, left, right) == left : F32}
87
88# ReluBackward is PyTorch's relu gradient, threshold_backward(g, relu(x), 0),
89# with x on the left and g on the right: +0.0 where x is at most zero, so the
90# gradient at +0.0 and -0.0 is +0.0, and g's own word otherwise, NaN x
91# included.
92law scalar_relu_backward_zeroed:
93  for +left: F32
94  for +right: F32
95  for zeroed: {Elementwise.at_most_zero(left) == True{} : Bool}
96  {Elementwise.scalar(Elementwise.ReluBackward{}, left, right) == 0.0 : F32}
97
98law scalar_relu_backward_passed:
99  for +left: F32
100  for +right: F32
101  for passed: {Elementwise.at_most_zero(left) == False{} : Bool}
102  {Elementwise.scalar(Elementwise.ReluBackward{}, left, right) == right : F32}
103
104# The words PyTorch 2.14.1's torch.relu returns at the edge words.
105law relu_edges:
106  {scalar_at(Elementwise.ReluLeft{}, 0, edges()) ==
107    [0, 2147483648, 1, 0, 8388607, 0, 8388608, 0, 1065353216, 0,
108      2139095039, 0, 2139095040, 0, 2139095041, 4286578689, 2143289344, 4290772992,
109      2147483647, 4294967295] : List<&2, U32>}
110
111# The gradient words PyTorch 2.14.1 gives torch.relu at the edge words, for
112# g = 2.5 (0x40200000).
113law relu_backward_edges:
114  {scalar_at(Elementwise.ReluBackward{}, 1075838976, edges()) ==
115    [0, 0, 1075838976, 0, 1075838976, 0, 1075838976, 0, 1075838976, 0,
116      1075838976, 0, 1075838976, 0, 1075838976, 1075838976, 1075838976, 1075838976,
117      1075838976, 1075838976] : List<&2, U32>}