More numerically stable jax.experimental.jet._integer_pow_taylor.

The previous version of `jax.experimental.jet._integer_pow_taylor` became numerically unstable when we evaluated the derivative jet of `x**n` around `x = 0` for integer `n > 2`. This was especially true in the higher order derivatives in `float32`.

Beforehand there was a complicated expression for approximating the jet of `x**n`. In this CL we replace integer powers with explicit multiplication, e.g. `x**2` becomes `x*x`. For higher powers of `x`, we use the exponentiation by squaring trick, which evaluates expressions such as `x**4` as `(x * x) * (x * x)`, which uses log(n) multiplies instead of `n-1`.

PiperOrigin-RevId: 617205359

Author: jax authors

jax/experimental/jet.py

@@ -490,24 +490,22 @@ def _pow_taylor(primals_in, series_in):
   return primal_out, series_out
 jet_rules[lax.pow_p] = _pow_taylor
 
+def _pow_by_squaring(x, n):
+  if n < 0:
+    return _pow_by_squaring(1 / x, -n)
+  elif n == 0:
+    return 1
+  elif n % 2 == 0:
+    return _pow_by_squaring(x * x, n / 2)
+  elif n % 2 == 1:
+    return x * _pow_by_squaring(x * x, (n - 1) / 2)
+
 def _integer_pow_taylor(primals_in, series_in, *, y):
   if y == 0:
     return jet(jnp.ones_like, primals_in, series_in)
-  elif y == 1:
-    return jet(lambda x: x, primals_in, series_in)
-  elif y == 2:
-    return jet(lambda x: x * x, primals_in, series_in)
-  x, = primals_in
-  series, = series_in
-  u = [x] + series
-  v = [lax.integer_pow(x, y)] + [None] * len(series)
-  for k in range(1, len(v)):
-    vu = sum(_scale(k, j) * v[k-j] * u[j] for j in range(1, k + 1))
-    uv = sum(_scale(k, j) * u[k-j] * v[j] for j in range(1, k))
-    v[k] = jnp.where(x == 0, 0, fact(k-1) * (y * vu - uv) / x)
-  primal_out, *series_out = v
+  else:
+    return jet(lambda x: _pow_by_squaring(x, y), primals_in, series_in)
 
-  return primal_out, series_out
 jet_rules[lax.integer_pow_p] = _integer_pow_taylor