math_fp32.bend
stelliferous@0.0.3.0 · 412 lines · 0 definitions
1# Ordered FP32 arithmetic on the shared sequence model. No separate container.
2import Base
3import stelliferous@0.0.3.0/math_sequence.bend as Sequence
4
5def multiply_add(accumulator: F32,input: F32,weight: F32) -> F32:
6 (accumulator + input * weight : F32)
7
8def weighted_add(weight: F32,accumulator: F32,input: F32) -> F32:
9 multiply_add(accumulator,input,weight)
10
11def zip_mac(length: Nat,accumulator: Sequence.Values(F32,length),input: Sequence.Values(F32,length),weight: F32) -> Sequence.Values(F32,length):
12 Sequence.zip_map(~F32,~F32,~weighted_add,length,weight,accumulator,input)
13
14# The Cephes expf polynomial after its constant and linear terms.
15def exponential_polynomial(+r: F32) -> F32:
16 (((((0.00019875691500 * r + 0.0013981999507 : F32) * r + 0.0083334519073 : F32) * r + 0.041665795894 : F32) * r + 0.16666665459 : F32) * r + 0.50000001201 : F32)
17
18# The 32-bit word of a U32, read as an F32, and back.
19def from_bits(value: U32) -> F32:
20 U32{word} = value
21 F32{word}
22
23def to_bits(value: F32) -> U32:
24 F32{bits} = value
25 U32{bits}
26
27# exp(r)-1-r near |r| <= ln(2)/2. Chebyshev coefficients, rounded to F32;
28# paired evaluation shortens the dependency chain without contracting products.
29def exponential_remainder(+r: F32) -> F32:
30 +z = (r * r : F32)
31 (z * ((0.5 + 0.16666577756404877 * r : F32) + z * ((0.041666556149721146 + 0.008363173343241215 * r : F32) + z * 0.0013926175888627768 : F32) : F32) : F32)
32
33# b=RN(1+q). Recover the small term lost in that sum before adding the
34# reduced exponential. Writing (1-b)+q keeps the eight lanes together under
35# the default vectorizer; q-(b-1) makes it mix unrelated addition trees.
36def logistic_denominator(q: F32,+base: F32,r: F32,correction: F32) -> F32:
37 (base + (r + (correction + ((1.0 - base : F32) + q : F32) : F32) : F32) : F32)
38
39def selected(+mask: U32,yes: F32,+no: F32) -> F32:
40 from_bits(U32.xor(F32.bits(no),U32.and(U32.xor(F32.bits(yes),F32.bits(no)),mask)))
41
42# Above 18 SiLU rounds to x. Pass x through, including positive infinity;
43# the negative tail uses the clamped input before the final underflow.
44def silu_finish(+x: F32,value: F32) -> F32:
45 +word = F32.bits(x)
46 +positive = U32.not((0 - U32.shrn(word,31n) : U32))
47 +large = U32.and(positive,(0 - U32.shrn((1099956224 - word : U32),31n) : U32))
48 selected(large,x,value)
49
50def logistic_scaled(silu: Bool,x: F32,value: F32,power: F32,denominator: F32) -> F32:
51 match silu:
52 case True{}:
53 silu_finish(x,(((value * power : F32) / denominator : F32) * 0.00000000000000000005421010862427522 : F32))
54 case False{}:
55 ((power / denominator : F32) * 0.00000000000000000005421010862427522 : F32)
56
57# u=-value, n approximates u/ln(2) to the nearest integer, r=u-n*ln(2).
58# The exponent word makes
59# H=2^(64-n), so sigmoid(x)=H/(exp(r)+2^-n)*2^-64. H remains normal;
60# only the last multiplication underflows. The tiny denominator addend q
61# may be floored to the least normal F32 without changing its rounded sum.
62def sigmoid_reduced(silu: Bool,x: F32,+value: F32,+biased: F32) -> F32:
63 +n = (biased - 12582912.0 : F32)
64 +r = ((n * F32.neg(0.693359375) - value : F32) + n * 0.00021219444005469058 : F32)
65 +power = from_bits((1602224128 - U32.shln(F32.bits(biased),23n) : U32))
66 +q = from_bits((U32.max(F32.bits(power),545259520) - 536870912 : U32))
67 logistic_scaled(silu,x,value,power,logistic_denominator(q,(1.0 + q : F32),r,exponential_remainder(r)))
68
69def sigmoid_input(silu: Bool,x: F32,+value: F32) -> F32:
70 sigmoid_reduced(silu,x,value,(value * F32.neg(1.44269504088896341) + 12582912.0 : F32))
71
72# At x>=18 sigmoid rounds to 1. Below -110 both sigmoid and SiLU round
73# to zero; multiplying the clamped x before the final scaling preserves -0
74# for SiLU, including at -infinity. NaNs pass the clamp and propagate.
75def sigmoid(+x: F32) -> F32:
76 sigmoid_input(False{},x,F32.max(F32.min(18.0,x),F32.neg(110.0)))
77
78def activate(+x: F32,enabled: Bool) -> F32:
79 match enabled:
80 case True{}:
81 sigmoid_input(True{},x,F32.max(F32.min(18.0,x),F32.neg(110.0)))
82 case False{}:
83 x
84
85# Twelve-bit heads have exact normal head products; the residual carries
86# the remaining twelve fraction bits.
87def split_high(+a: F32) -> F32:
88 from_bits(U32.and(F32.bits(a),4294963200))
89
90# Chebyshev interpolation of (exp(r)-1-r)/r^2 on [-0.375,1.0], rounded to
91# F32. The compensated head needs a wider interval than ln(2)/2. Paired
92# evaluation shortens the dependency chain without contracting products.
93def gelu_exponential_remainder(+r: F32) -> F32:
94 +z = (r * r : F32)
95 (z * ((0.5 + 0.16666662693023682 * r : F32) + z * ((0.04166663810610771 + 0.008334293961524963 * r : F32) + z * ((0.0013882070779800415 + 0.00019411252287682146 * r : F32) + z * 3.194776218151674e-05 : F32) : F32) : F32) : F32)
96
97def gelu_reduced(x: F32,+value: F32,low: F32,+biased: F32) -> F32:
98 +n = (biased - 12582912.0 : F32)
99 +r = (((n * F32.neg(0.693359375) - value : F32) + n * 0.00021219444005469058 : F32) - low : F32)
100 +power = from_bits((1602224128 - U32.shln(F32.bits(biased),23n) : U32))
101 +q = from_bits((U32.max(F32.bits(power),545259520) - 536870912 : U32))
102 (((x * power : F32) / logistic_denominator(q,(1.0 + q : F32),r,gelu_exponential_remainder(r)) : F32) * 0.00000000000000000005421010862427522 : F32)
103
104# w=2 sqrt(2/pi) c (1+0.044715 c^2). Twelve-bit input/product heads,
105# a nine-bit cube head and short coefficient heads keep the head sum exact
106# for 0.5<=|c|<=16: the linear/cubic products need at most 18/16 bits,
107# and their exponent gap is at most 7/4 bits. Smaller inputs have a small
108# absolute head error, measured separately from the ordered-map theorem.
109# Carry coefficient and cubic tails through c_low^2, omitting c_low^3.
110type Argument is Data:
111 Argument{x: F32,head: F32,low: F32}
112
113def gelu_argument(+x: F32,+c: F32) -> Argument:
114 +h = split_high(c)
115 +l = (c - h : F32)
116 +p = (h * h : F32)
117 +ph = split_high(p)
118 +cube = (h * ph : F32)
119 +cubeh = from_bits(U32.and(F32.bits(cube),4294934528))
120 +linear = (1.59375 * h : F32)
121 +cubic = (0.0712890625 * cubeh : F32)
122 +w = (linear + cubic : F32)
123 low = (((0.0020191215444356203 * h : F32) +
124 ((0.0712890625 * (cube - cubeh : F32) : F32) + (6.575377483386546e-05 * cube : F32) : F32) : F32) +
125 (((0.07135481387376785 * h : F32) * (p - ph : F32) : F32) +
126 (l * (1.5957691669464111 + 0.21406444907188416 * (p + h * l : F32) : F32) : F32) : F32) : F32)
127 +value = F32.max(w,F32.neg(110.0))
128 Argument{x,value,low}
129
130# Positive results already round to x at magnitude >=6. Cap their
131# negative-tail evaluation there to avoid an unnecessary final underflow;
132# negative results keep the full range through magnitude 16.
133def gelu_negative_input(+x: F32) -> F32:
134 +word = F32.bits(x)
135 +negative = (0 - U32.shrn(word,31n) : U32)
136 F32.max(from_bits(U32.or(word,2147483648)),selected(negative,F32.neg(16.0),F32.neg(6.0)))
137
138def gelu_finish(arg: Argument) -> F32:
139 Argument{+x,+value,low} = arg
140 +negative_input = gelu_negative_input(x)
141 +tail = gelu_reduced(negative_input,value,low,(value * F32.neg(1.44269504088896341) + 12582912.0 : F32))
142 selected((0 - U32.shrn(F32.bits(x),31n) : U32),tail,(x + tail : F32))
143
144# GELU(x)=x+GELU(-x). Evaluate a negative tail on both sides, then
145# restore the positive result. Clamping before arithmetic keeps infinities
146# out of the products. Negative magnitudes >=16 round to -0.
147# The original word selects the sign, including -0; NaNs propagate.
148def gelu_prepared(+x: F32) -> Argument:
149 gelu_argument(x,gelu_negative_input(x))
150
151def gelu_tanh(x: F32) -> F32:
152 gelu_finish(gelu_prepared(x))
153
154# Sign restoration and the tiny interval are word selections. For |x|<2^-12,
155# |x-tanh(x)|<|x|^3/3 is below half an F32 ulp, so return x, including -0.
156def tanh_finish(+x: F32,+value: F32) -> F32:
157 +word = F32.bits(x)
158 +small = (0 - U32.shrn((U32.and(word,2147483647) - 964689920 : U32),31n) : U32)
159 +keep = U32.or(small,2147483648)
160 from_bits(U32.or(U32.and(word,keep),U32.and(F32.bits(value),U32.not(keep))))
161
162# At |x|>=10, 1-tanh(|x|)<2*exp(-20)<2^-25. Clamping before arithmetic
163# also keeps infinities away from intermediate products; NaNs propagate.
164def tanh_magnitude(+x: F32) -> F32:
165 F32.max(F32.min(10.0,from_bits(U32.and(F32.bits(x),2147483647))),0.000244140625)
166
167# Chebyshev interpolation of (tanh(sqrt(z))/sqrt(z)-1)/z, |t|<=0.75.
168# The coefficients are F32 values.
169def tanh_small(+t: F32,+z: F32) -> F32:
170 (t + (t * z : F32) * (((((0.001909503829665482 * z + F32.neg(0.007933800108730793) : F32) * z + 0.021615320816636086 : F32) * z + F32.neg(0.05393586680293083) : F32) * z + 0.1333317905664444 : F32) * z + F32.neg(0.3333333134651184) : F32) : F32)
171
172# One additional term for the expm1 correction in the tanh quotient.
173def tanh_remainder(+r: F32) -> F32:
174 ((((((0.00019890980911441147 * r + 0.0013933641603216529 : F32) * r + 0.00833331048488617 : F32) * r + 0.04166646674275398 : F32) * r + 0.1666666716337204 : F32) * r + 0.5 : F32) * (r * r : F32) : F32)
175
176def tanh_exponential(+a: F32,+biased: F32,x: F32) -> F32:
177 +n = (biased - 12582912.0 : F32)
178 +r = (((a * 2.0 : F32) - n * 0.693359375 : F32) - n * F32.neg(0.00021219444005469058) : F32)
179 +q = from_bits((1065353216 - U32.shln(F32.bits(biased),23n) : U32))
180 +t = F32.min(0.75,a)
181 +small = tanh_small(t,(t * t : F32))
182 +large = (1.0 - 2.0 * (q / ((1.0 + q : F32) + (r + tanh_remainder(r) : F32) : F32) : F32) : F32)
183 +mask = (0 - U32.shrn((F32.bits(a) - 1061158912 : U32),31n) : U32)
184 tanh_finish(x,selected(mask,small,large))
185
186def tanh_input(+a: F32,x: F32) -> F32:
187 tanh_exponential(a,((a * 2.0 : F32) * 1.44269504088896341 + 12582912.0 : F32),x)
188
189def tanh(+x: F32) -> F32:
190 tanh_input(tanh_magnitude(x),x)
191
192# Split 2^n at +64 for nonnegative inputs and -64 for negative inputs.
193# The first product is always normal and exact; only the last product can
194# underflow. The exponent words come directly from the biased integer n.
195def exponential_scale(biased: F32,x: F32,p: F32) -> F32:
196 +sign = U32.shrn(U32.and(F32.bits(x),2147483648),1n)
197 ((p * from_bits((1602224128 - sign : U32)) : F32) *
198 from_bits((528482304 + U32.shln(F32.bits(biased),23n) + sign : U32)) : F32)
199
200# FastTwoSum retains the low bits of 1+r before adding the quadratic tail.
201def exponential_precise(+r: F32) -> F32:
202 +sum = (1.0 + r : F32)
203 +error = (r - (sum - 1.0 : F32) : F32)
204 (sum + ((exponential_polynomial(r) * (r * r : F32) : F32) + error : F32) : F32)
205
206def exponential_reduced(+x: F32,+biased: F32) -> F32:
207 +n = (biased - 12582912.0 : F32)
208 +high = (x - n * 0.693359375 : F32)
209 +low = (n * F32.neg(0.00021219444005469058) : F32)
210 +r = (high - low : F32)
211 exponential_scale(biased,x,exponential_precise(r))
212
213def exponential_input(+x: F32) -> F32:
214 exponential_reduced(x,(x * 1.44269504088896341 + 12582912.0 : F32))
215
216def exp(x: F32) -> F32:
217 exponential_input(F32.max(F32.min(89.0,x),F32.neg(104.0)))
218
219# For s near r/(2+r), d = r-(2+r)*s gives the exact real identity
220# log(1+r) = 2*atanh(s) + log(1+d/(1+s)). The odd series ends at s^11.
221# Split r and s into 12-bit heads so their head product is exact; retain
222# the low products in d. (1-s)*(1+s^2+s^4) approximates 1/(1+s).
223# FastTwoSum retains the low bits when adding the exponent's ln(2) part.
224def logarithm_series(+r: F32,+n: F32) -> F32:
225 +s = (r / (2.0 + r : F32) : F32)
226 +square = (s * s : F32)
227 +r_high = from_bits(U32.and(F32.bits(r),4294963200))
228 +r_low = (r - r_high : F32)
229 +s_high = from_bits(U32.and(F32.bits(s),4294963200))
230 +s_low = (s - s_high : F32)
231 +product_low = ((r_high * s_low + r_low * s_high : F32) + r_low * s_low : F32)
232 +correction = (((r - 2.0 * s : F32) - r_high * s_high : F32) - product_low : F32)
233 +p = ((((0.09090909090909091 * square + 0.1111111111111111 : F32) * square + 0.14285714285714285 : F32) * square + 0.2 : F32) * square + 0.3333333333333333 : F32)
234 +main = (2.0 * s : F32)
235 +rest = ((main * square : F32) * p + correction * ((1.0 - s : F32) * (1.0 + square * (1.0 + square : F32) : F32) : F32) : F32)
236 +high = (n * 0.693359375 : F32)
237 +sum = (high + main : F32)
238 +tail = (main - (sum - high : F32) : F32)
239 (sum + (rest + (n * F32.neg(0.000212194440) + tail : F32) : F32) : F32)
240
241# The exponent carry at mantissa word 0x3504f4 selects the same reduced
242# interval as halving m above sqrt(2), without a floating-point selection.
243def logarithm_normalized(x: F32,adjustment: F32) -> F32:
244 +word = F32.bits(x)
245 +adjusted = (word + 4913932 : U32)
246 logarithm_series((from_bits((word - U32.and(adjusted,4286578688) + 1065353216 : U32)) - 1.0 : F32),
247 ((from_bits(U32.or(U32.and(U32.shrn(adjusted,23n),255),1258291200)) - 8388735.0 : F32) - adjustment : F32))
248
249# Scaling a subnormal by 2^23 is exact. Zero, negatives, infinities and
250# NaNs are selected explicitly, after normalizing the finite positive input.
251def logarithm_finish(+x: F32,value: F32) -> F32:
252 +word = F32.bits(x)
253 +magnitude = U32.and(word,2147483647)
254 +zero = (0 - U32.shrn((magnitude - 1 : U32),31n) : U32)
255 +negative = (0 - U32.shrn(word,31n) : U32)
256 +special = U32.not((0 - U32.shrn((magnitude - 2139095040 : U32),31n) : U32))
257 selected(zero,from_bits(4286578688),selected(negative,from_bits(2143289344),selected(special,(x + x : F32),value)))
258
259def logarithm_input(+x: F32) -> F32:
260 +small = (0 - U32.shrn((U32.and(F32.bits(x),2147483647) - 8388608 : U32),31n) : U32)
261 logarithm_finish(x,logarithm_normalized(selected(small,(x * 8388608.0 : F32),x),from_bits(U32.and(small,1102577664))))
262
263def log(x: F32) -> F32:
264 logarithm_input(x)
265
266# The square root rounded to nearest, as IEEE 754 requires of every backend.
267def sqrt(x: F32) -> F32:
268 F32.sqrt(x)
269
270# Cancellation-free even derivatives with one exponential reduction.
271def sigmoid_derivative(x: F32) -> F32:
272 +q = exp(F32.neg(from_bits(U32.and(F32.bits(x),2147483647))))
273 (q / (1.0 + q * (2.0 + q : F32) : F32) : F32)
274
275def tanh_derivative(x: F32) -> F32:
276 +q = exp((F32.neg(2.0) * from_bits(U32.and(F32.bits(x),2147483647)) : F32))
277 ((4.0 * q : F32) / (1.0 + q * (2.0 + q : F32) : F32) : F32)
278
279# The derivative divided by the distance from its zero, a degree-eight
280# polynomial on [-2,-0.5] in the exact difference from the head of the zero.
281def silu_derivative_root(+value: F32) -> F32:
282 +d = (value + 1.2784645557403564 : F32)
283 +z = (d * d : F32)
284 +p0 = (0.2178117036819458 + 0.14664886891841888 * d : F32)
285 +p1 = (0.01887473464012146 + F32.neg(0.01522465143352747) * d : F32)
286 +p2 = (F32.neg(0.006606101058423519) + 0.0001401008339598775 * d : F32)
287 +p3 = (0.0007981295930221677 + 0.00015176329179666936 * d : F32)
288 +tail = (p0 + z * (p1 + z * (p2 + z * (p3 + z * F32.neg(4.2798692447831854e-05) : F32) : F32) : F32) : F32)
289 ((d - 1.2979282537628478e-08 : F32) * tail : F32)
290
291def silu_derivative_tail(+value: F32,+scaled: F32) -> F32:
292 +s = (scaled * 0.00000000000000000005421010862427522 : F32)
293 +word = F32.bits(value)
294 +above = U32.not((0 - U32.shrn((word - 3204448256 : U32),31n) : U32))
295 +below = U32.not((0 - U32.shrn((3221225472 - word : U32),31n) : U32))
296 +tail = ((scaled * (1.0 + (value - value * s : F32) : F32) : F32) * 0.00000000000000000005421010862427522 : F32)
297 selected(U32.and(above,below),silu_derivative_root(value),tail)
298
299type DerivativeArgument is Data:
300 DerivativeArgument{x: F32,value: F32,scaled: F32}
301
302def silu_derivative_prepared(+x: F32,+value: F32) -> DerivativeArgument:
303 +biased = (value * F32.neg(1.44269504088896341) + 12582912.0 : F32)
304 +n = (biased - 12582912.0 : F32)
305 +r = ((n * F32.neg(0.693359375) - value : F32) + n * 0.00021219444005469058 : F32)
306 +power = from_bits((1602224128 - U32.shln(F32.bits(biased),23n) : U32))
307 +q = from_bits((U32.max(F32.bits(power),545259520) - 536870912 : U32))
308 DerivativeArgument{x,value,(power / logistic_denominator(q,(1.0 + q : F32),r,exponential_remainder(r)) : F32)}
309
310def silu_derivative_prepare(+x: F32) -> DerivativeArgument:
311 silu_derivative_prepared(x,F32.max(from_bits(U32.or(F32.bits(x),2147483648)),F32.neg(110.0)))
312
313def silu_derivative_finish(arg: DerivativeArgument) -> F32:
314 DerivativeArgument{+x,value,scaled} = arg
315 +tail = silu_derivative_tail(value,scaled)
316 +magnitude = U32.and(F32.bits(x),2147483647)
317 +nan = (0 - U32.shrn((2139095040 - magnitude : U32),31n) : U32)
318 selected(nan,x,selected((0 - U32.shrn(F32.bits(x),31n) : U32),tail,(1.0 - tail : F32)))
319
320def silu_derivative(x: F32) -> F32:
321 silu_derivative_finish(silu_derivative_prepare(x))
322
323# Root quotient on [-1.25,-0.5]. Split the linear head exactly and
324# retain the low parts of the first two rounded quotient coefficients.
325def gelu_derivative_root(+c: F32) -> F32:
326 +d = (c + 0.7524614334106445 : F32)
327 +dh = from_bits(U32.and(F32.bits(d),4294967292))
328 +dl = (d - dh : F32)
329 +z = (d * d : F32)
330 +p0 = (0.05540008842945099 + 0.38751843571662903 * d : F32)
331 +p1 = (F32.neg(0.01578284427523613) + F32.neg(0.11394448578357697) * d : F32)
332 +p2 = (F32.neg(0.01661972515285015) + 0.01968197152018547 * d : F32)
333 +p3 = (0.005269420798867941 + F32.neg(0.0024054001551121473) * d : F32)
334 +p4 = (F32.neg(0.0009768777526915073) + 8.785235695540905e-05 * d : F32)
335 +tail = (p0 + z * (p1 + z * (p2 + z * (p3 + z * p4 : F32) : F32) : F32) : F32)
336 ((0.375 * dh : F32) + (((d - 1.1339628272762046e-08 : F32) * tail : F32) + (((0.375 * dl : F32) - 4.25236024170772e-09 : F32) + d * (F32.neg(1.8141816937600197e-09) + d * 1.0859928956108433e-08 : F32) : F32) : F32) : F32)
337
338type GeluPrepared is Data:
339 GeluPrepared{x: F32,head: F32,low: F32,slope: F32,slope_error: F32}
340
341def gelu_derivative_prepared(+c: F32,arg: Argument) -> GeluPrepared:
342 Argument{x,head,+low} = arg
343 +h = split_high(c)
344 +l = (c - h : F32)
345 +cube = (h * split_high((h * h : F32)) : F32)
346 +cubeh = from_bits(U32.and(F32.bits(cube),4294934528))
347 +slope = (1.59375 * h + 3.0 * (0.0712890625 * cubeh : F32) : F32)
348 +slope_error = ((3.0 * low : F32) - ((0.004038243088871241 * h : F32) + (3.1915383338928223 * l : F32) : F32) : F32)
349 GeluPrepared{x,head,low,slope,slope_error}
350
351def gelu_derivative_prepare(+x: F32) -> GeluPrepared:
352 +c = F32.max(from_bits(U32.or(F32.bits(x),2147483648)),F32.neg(11.0))
353 gelu_derivative_prepared(c,gelu_argument(x,c))
354
355type GeluReduced is Data:
356 GeluReduced{x: F32,power: F32,denominator: F32,error: F32,slope: F32,slope_error: F32}
357
358def gelu_derivative_exponential(arg: GeluPrepared) -> GeluReduced:
359 GeluPrepared{x,+head,low,slope,slope_error} = arg
360 +biased = (head * F32.neg(1.44269504088896341) + 12582912.0 : F32)
361 +n = (biased - 12582912.0 : F32)
362 +rh = (n * F32.neg(0.693359375) - head : F32)
363 +rl = (n * 0.00021219444005469058 - low : F32)
364 +r = (rh + rl : F32)
365 +v = (r - rh : F32)
366 +rerror = ((rh - (r - v : F32) : F32) + (rl - v : F32) : F32)
367 +power = from_bits((1602224128 - U32.shln(F32.bits(biased),23n) : U32))
368 +small = from_bits((U32.max(F32.bits(power),545259520) - 536870912 : U32))
369 +base = (1.0 + small : F32)
370 +tail = (r + (gelu_exponential_remainder(r) + ((1.0 - base : F32) + small : F32) : F32) : F32)
371 +denominator = (base + tail : F32)
372 +error = (((base - denominator : F32) + tail : F32) + (denominator - small : F32) * rerror : F32)
373 GeluReduced{x,power,denominator,error,slope,slope_error}
374
375# Correct the quotient before the final power-of-two scaling.
376def gelu_derivative_evaluated(arg: GeluReduced) -> DerivativeArgument:
377 GeluReduced{+x,+power,+denominator,error,+slope,+slope_error} = arg
378 +sigmoid = ((power / denominator : F32) * 0.00000000000000000005421010862427522 : F32)
379 +factor = (1.0 + slope : F32)
380 +residual = (slope_error - (slope + slope_error : F32) * sigmoid : F32)
381 +quotient = split_high((factor / denominator : F32))
382 +dh = split_high(denominator)
383 +dl = (denominator - dh : F32)
384 +remainder = (((factor - quotient * dh : F32) - quotient * dl : F32) + (residual - quotient * error : F32) : F32)
385 +value = (quotient + remainder / denominator : F32)
386 +result = ((power * value : F32) * 0.00000000000000000005421010862427522 : F32)
387 DerivativeArgument{x,F32.max(from_bits(U32.or(F32.bits(x),2147483648)),F32.neg(11.0)),result}
388
389# Split the linear term before adding the odd correction to one half.
390def gelu_derivative_small(+x: F32) -> F32:
391 +h = split_high(x)
392 +l = (x - h : F32)
393 +linear = (0.796875 * h : F32)
394 +base = (0.5 + linear : F32)
395 +z = (x * x : F32)
396 +p = (F32.neg(0.26727786660194397) + z * (0.061208296567201614 + z * (F32.neg(0.009932546876370907) + z * 0.0011246189242228866 : F32) : F32) : F32)
397 +low = ((0.5 - base : F32) + linear : F32)
398 (base + (low + ((0.0010095608028653559 * h + 0.7978845834732056 * l : F32) + (x * z : F32) * p : F32) : F32) : F32)
399
400def gelu_derivative_finish(arg: DerivativeArgument) -> F32:
401 DerivativeArgument{+x,+value,result} = arg
402 +word = F32.bits(value)
403 +above = U32.not((0 - U32.shrn((word - 3204448256 : U32),31n) : U32))
404 +below = U32.not((0 - U32.shrn((3214934016 - word : U32),31n) : U32))
405 +tail = selected(U32.and(above,below),gelu_derivative_root(value),result)
406 +magnitude = U32.and(F32.bits(x),2147483647)
407 +nan = (0 - U32.shrn((2139095040 - magnitude : U32),31n) : U32)
408 +small = (0 - U32.shrn((magnitude - 1056964608 : U32),31n) : U32)
409 selected(nan,x,selected(small,gelu_derivative_small(x),selected((0 - U32.shrn(F32.bits(x),31n) : U32),tail,(1.0 - tail : F32))))
410
411def gelu_tanh_derivative(x: F32) -> F32:
412 gelu_derivative_finish(gelu_derivative_evaluated(gelu_derivative_exponential(gelu_derivative_prepare(x))))