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

Merge pull request #20923 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
jax/_src/test_util.py
View File

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