File size: 4,736 Bytes
073de96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """erf in float64 over NumPy arrays (NumPy has none): FreeBSD msun s_erf.c as ported by musl and Rust `libm`,
vectorized region by region. Used by the head's GELU (the author's torch.nn.GELU / F.gelu, the exact erf form).
* ====================================================
* Copyright (C) 1993 by Sun Microsystems, Inc. All rights reserved.
*
* Developed at SunPro, a Sun Microsystems, Inc. business.
* Permission to use, copy, modify, and distribute this
* software is freely granted, provided that this notice
* is preserved.
* ===================================================="""
from __future__ import annotations
import math
import numpy as np
ERX = 8.45062911510467529297e-01
EFX8 = 1.02703333676410069053e+00
PP = (1.28379167095512558561e-01, -3.25042107247001499370e-01, -2.84817495755985104766e-02,
-5.77027029648944159157e-03, -2.37630166566501626084e-05)
QQ = (3.97917223959155352819e-01, 6.50222499887672944485e-02, 5.08130628187576562776e-03,
1.32494738004321644526e-04, -3.96022827877536812320e-06)
PA = (-2.36211856075265944077e-03, 4.14856118683748331666e-01, -3.72207876035701323847e-01,
3.18346619901161753674e-01, -1.10894694282396677476e-01, 3.54783043256182359371e-02,
-2.16637559486879084300e-03)
QA = (1.06420880400844228286e-01, 5.40397917702171048937e-01, 7.18286544141962662868e-02,
1.26171219808761642112e-01, 1.36370839120290507362e-02, 1.19844998467991074170e-02)
RA = (-9.86494403484714822705e-03, -6.93858572707181764372e-01, -1.05586262253232909814e+01,
-6.23753324503260060396e+01, -1.62396669462573470355e+02, -1.84605092906711035994e+02,
-8.12874355063065934246e+01, -9.81432934416914548592e+00)
SA = (1.96512716674392571292e+01, 1.37657754143519042600e+02, 4.34565877475229228821e+02,
6.45387271733267880336e+02, 4.29008140027567833386e+02, 1.08635005541779435134e+02,
6.57024977031928170135e+00, -6.04244152148580987438e-02)
RB = (-9.86494292470009928597e-03, -7.99283237680523006574e-01, -1.77579549177547519889e+01,
-1.60636384855821916062e+02, -6.37566443368389627722e+02, -1.02509513161107724954e+03,
-4.83519191608651397019e+02)
SB = (3.03380607434824582924e+01, 3.25792512996573918826e+02, 1.53672958608443695994e+03,
3.19985821950859553908e+03, 2.55305040643316442583e+03, 4.74528541206955367215e+02,
-2.24409524465858183362e+01)
SQRT1_2 = math.sqrt(0.5)
def _poly(z, c):
"""c[0] + z * (c[1] + z * (... + z * c[-1])), Horner from the last coefficient (the C source's nesting)."""
r = np.full_like(z, c[-1])
for a in c[-2::-1]:
r = a + z * r
return r
def _one_plus_z_poly(z, c):
"""1 + z * (c[0] + z * (c[1] + ...)) (the denominators Q / S)."""
return 1.0 + z * _poly(z, c)
def erf(x) -> np.ndarray:
"""Elementwise erf of a float64 array (NaN -> NaN, +-inf -> +-1)."""
x = np.asarray(x, dtype=np.float64)
out = np.empty_like(x)
ax = np.abs(x)
hi = (ax.view(np.uint64) >> np.uint64(32)).astype(np.uint32) # high word of |x|
sign = np.signbit(x)
m = hi < 0x3e300000 # |x| < 2**-28
out[m] = 0.125 * (8.0 * x[m] + EFX8 * x[m])
m = (hi >= 0x3e300000) & (hi < 0x3feb0000) # |x| < 0.84375
xm = x[m]
z = xm * xm
out[m] = xm + xm * (_poly(z, PP) / _one_plus_z_poly(z, QQ))
m = (hi >= 0x3feb0000) & (hi < 0x3ff40000) # 0.84375 <= |x| < 1.25
s = ax[m] - 1.0
erfc1 = 1.0 - ERX - _poly(s, PA) / _one_plus_z_poly(s, QA)
out[m] = np.where(sign[m], -(1.0 - erfc1), 1.0 - erfc1)
m = (hi >= 0x3ff40000) & (hi < 0x40180000) # 1.25 <= |x| < 6
if m.any():
a = ax[m]
h = hi[m]
s = 1.0 / (a * a)
near = h < 0x4006db6d # |x| < 1 / 0.35
r = np.where(near, _poly(s, RA), _poly(s, RB))
big_s = np.where(near, _one_plus_z_poly(s, SA), _one_plus_z_poly(s, SB))
zz = (a.view(np.uint64) & np.uint64(0xFFFFFFFF00000000)).view(np.float64) # low word cleared
erfc2 = np.exp(-zz * zz - 0.5625) * np.exp((zz - a) * (zz + a) + r / big_s) / a
out[m] = np.where(sign[m], -(1.0 - erfc2), 1.0 - erfc2)
m = hi >= 0x40180000 # |x| >= 6, inf, NaN
y = 1.0 - math.ldexp(1.0, -1022)
out[m] = np.where(np.isnan(x[m]), np.nan, np.where(sign[m], -y, y))
return out
def gelu(x: np.ndarray) -> np.ndarray:
"""GELU, the exact form x / 2 * (1 + erf(x / sqrt 2)), evaluated in float64 and rounded to x's float32."""
xd = np.asarray(x, dtype=np.float64)
return (0.5 * xd * (1.0 + erf(xd * SQRT1_2))).astype(np.float32)
|