mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-26 06:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request from pearu:pearu/asinh-2

Browse Source 
PiperOrigin-RevId: 629130313
This commit is contained in:
jax authors 2024-04-29 11:31:59 -07:00
commit b44e9bfe66
1 changed files with 27 additions and 10 deletions

37
jax/_src/test_util.py

@ -1796,12 +1796,12 @@ class numpy_with_mpmath:
def arcsin(self, x): def arcsin(self, x):
ctx = x.context ctx = x.context
if isinstance(x, ctx.mpc): if isinstance(x, ctx.mpc):
# Workaround mpmath 1.3 bug in asin(+-inf+-infj) evaluation (see mpmath/mpmath#793). # Workaround mpmath 1.3 bug in asin(+-inf+-infj) evaluation (see
# TODO(pearu): remove this function when mpmath 1.4 or newer # mpmath/mpmath#793).
# will be the required test dependency. # TODO(pearu): remove the if-block below when mpmath 1.4 or
# newer will be the required test dependency.
pi = ctx.pi pi = ctx.pi
inf = ctx.inf inf = ctx.inf
nan = ctx.nan
zero = ctx.zero zero = ctx.zero
if ctx.isinf(x.real): if ctx.isinf(x.real):
sign_real = -1 if x.real < 0 else 1 sign_real = -1 if x.real < 0 else 1
@ -1811,9 +1811,10 @@ class numpy_with_mpmath:
elif ctx.isinf(x.imag): elif ctx.isinf(x.imag):
return ctx.make_mpc((zero._mpf_, x.imag._mpf_)) return ctx.make_mpc((zero._mpf_, x.imag._mpf_))
# TODO(pearu): adjust this code according to mpmath/mpmath#786 # On branch cut, mpmath.mp.asin returns different value compared
# resolution when mpmath 1.4 or newer will be the required test # to mpmath.fp.asin and numpy.arcsin (see
# dependency. # mpmath/mpmath#786). The following if-block ensures
# compatibiliy with numpy.arcsin.
if x.real > 1 and x.imag == 0: if x.real > 1 and x.imag == 0:
return ctx.asin(x).conjugate() return ctx.asin(x).conjugate()
@ -1822,10 +1823,26 @@ class numpy_with_mpmath:
def arcsinh(self, x): def arcsinh(self, x):
ctx = x.context ctx = x.context
# TODO(pearu): adjust this code according to mpmath/mpmath#786
# resolution when mpmath 1.4 or newer will be the required test
# dependency.
if isinstance(x, ctx.mpc): if isinstance(x, ctx.mpc):
# Workaround mpmath 1.3 bug in asinh(+-inf+-infj) evaluation
# (see mpmath/mpmath#749).
# TODO(pearu): remove the if-block below when mpmath 1.4 or
# newer will be the required test dependency.
pi = ctx.pi
inf = ctx.inf
zero = ctx.zero
if ctx.isinf(x.imag):
sign_imag = -1 if x.imag < 0 else 1
real = -inf if x.real < 0 else inf
imag = sign_imag * pi / (4 if ctx.isinf(x.real) else 2)
return ctx.make_mpc((real._mpf_, imag._mpf_))
elif ctx.isinf(x.real):
return ctx.make_mpc((x.real._mpf_, zero._mpf_))
# On branch cut, mpmath.mp.asinh returns different value
# compared to mpmath.fp.asinh and numpy.arcsinh (see
# mpmath/mpmath#786). The following if-block ensures
# compatibiliy with numpy.arcsinh.
if x.real == 0 and x.imag < -1: if x.real == 0 and x.imag < -1:
return (-ctx.asinh(x)).conjugate() return (-ctx.asinh(x)).conjugate()
return ctx.asinh(x) return ctx.asinh(x)