`jnp.arcsinh` / `jnp.arccosh` JVP is exactly 0 in float32 for $|x| > ~1.84e19$ because of squared x
Description
[!NOTE] The issue was identified by Claude. It came up during numerical stabilization work for a classification layer in hyperbolic deep learning. Concretely, a Multinomial Logistic Regression (MLR) head computes its score as $|z|\cdot \text{arcsinh} \big( \langle x, w \rangle_L \big)$, with arguments growing like $e^{radius}$, and $\langle \cdot \rangle_L$ being the Minkowski inner product.
Description
The JVP rules for asinh_p and acosh_p compute the derivative as $1/\sqrt{x^2 \pm 1}$. In float32, $x^2$ overflows to inf once $|x| \gtrsim 1.84 \times 10^{19}$, so rsqrt(inf) = 0 and the derivative is returned as exactly 0.0. The true derivative $1/\sqrt{1+x^2} \approx 1/|x|$ is still representable down to $|x| \approx 10^{38}$.
The forward value is correct at these arguments (the CHLO asinh/acosh lowerings use a log form). Only the derivative is wrong. In a chain $\text{arcsinh}(f(x))$ where $f$ grows exponentially, the true gradient is $f'/f$, which is $O(1)$. The overflow reduces it to $0 \cdot f' = 0$, so a parameter silently stops training. The issue occurs on both the CPU and the GPU.
Reproduction
import jax
import jax.numpy as jnp
jax.config.update('jax_enable_x64', True)
x = jnp.float32(1e22)
print("grad f32 :", jax.grad(jnp.arcsinh)(x))
print("grad f64 :", jax.grad(jnp.arcsinh)(jnp.float64(1e22))) # with jax_enable_x64
print("1/hypot f32:", 1 / jnp.hypot(jnp.float32(1), x))
print("value f32 :", jnp.arcsinh(x))
print("acosh grad :", jax.grad(jnp.arccosh)(x))
Expected output
grad f32 : 1e-22
grad f64 : 1e-22
1/hypot f32: 1e-22
value f32 : 51.350018
acosh grad : 1e-22
Actual output
grad f32 : 0.0
grad f64 : 1e-22
1/hypot f32: 1e-22
value f32 : 51.350018
acosh grad : 0.0
The float32 gradient is 1e-19 at x = 1e19 and exactly 0.0 from x = 3e19 on.
Root cause
square(x) overflows to inf before the result underflows, and rsqrt(inf) = 0. The comment above the acosh rule explains why the $x^2 - 1$ form was chosen over $(x+1)(x-1)$, but neither formula is safe at large $x$.
Suggested fix
asinh_p: replace rsqrt(square(x) + 1) with reciprocal(hypot(1, x)).
acosh_p: replace rsqrt(square(x) − 1) with reciprocal(sqrt(x − 1) · sqrt(x + 1)). The factored form avoids squaring $x$ and so cannot overflow, while remaining accurate on the domain $x\geq1$.
Alternatively, scale by $\max(1, |x|)$ before squaring in both.
System info (python version, jaxlib version, accelerator, etc.)
jax: 0.11.1
jaxlib: 0.11.1
numpy: 2.5.3
python: 3.12.7 (main, Oct 16 2024, 04:37:19) [Clang 18.1.8 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='rey', release='6.8.0-54-generic', version='#56-Ubuntu SMP PREEMPT_DYNAMIC Sat Feb 8 00:37:57 UTC 2025', machine='x86_64')
JAX_ENABLE_X64=1
$ nvidia-smi
Fri Sep 11 13:33:09 2026
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.216.03 Driver Version: 535.216.03 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA A100-PCIE-40GB Off | 00000000:5E:00.0 Off | 0 |
| N/A 37C P0 77W / 250W | 4711MiB / 40960MiB | 99% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 1 NVIDIA A100-PCIE-40GB Off | 00000000:86:00.0 Off | 0 |
| N/A 39C P0 87W / 250W | 6682MiB / 40960MiB | 91% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
Source: jax-ml/jax