lax.acosh: express gradient using x^2-1 directly.

Jake VanderPlas · Google-ML-Automation · commit fc24fb9bf529 · 2025-10-07T13:04:01.000-07:00
The (x+1)(x-1) form can be more precise as x-&gt;0, but since acosh is only defined for x&gt;=1, we can use the direct form. This saves one op and leads to cleaner second-order autodiff.

PiperOrigin-RevId: 816343493
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
@@ -4377,7 +4377,9 @@ def atan_impl(x):
 
 acosh_p = standard_unop(_float | _complex, 'acosh')
 ad.defjvp(acosh_p,
-          lambda g, x: mul(g, rsqrt(mul(sub(x, _one(x)), add(x, _one(x))))))
+          # We use x^2-1 rather than (x+1)(x-1). The latter is more accurate
+          # for x near zero, but the function domain is x>=1.
+          lambda g, x: mul(g, rsqrt(sub(square(x), _one(x)))))
 mlir.register_lowering(acosh_p, partial(_nary_lower_hlo, chlo.acosh))
 
 atanh_p = standard_unop(_float | _complex, 'atanh')
