Repository: jax-ml/jax
Pull request: Add a primitive integer_pow() for values raised to a fixed integer scalar.
Use `integer_pow()` in the RHS JVP of `div()`. Also use it in `square()` and `reciprocal()`.

Fixes #3136

```
In [1]: from jax import grad, make_jaxpr
In [2]: def inv(x): return 1/x
In [3]: print(grad(grad(grad(grad(grad(grad(inv))))))(4.))
0.043945312

In [4]: make_jaxpr(grad(grad(grad(grad(grad(grad(inv)))))))(4.)
Out[4]:
{ lambda  ; a.
  let b = integer_pow[ y=-7 ] a
      c = mul -6.0 b
      d = mul -120.0 c
  in (d,) }

In [5]:
```

Problem statement:
Higher order derivatives return wrong results
When trying to compute higher order derivatives of a very simple function:

```python
def inv(x):
    return 1/X

print(grad(grad(grad(grad(grad(grad(inv))))))(4.))
```

It prints `-1.9560547` instead of the correct `0.0439453125`. 
It starts returning incorrect results at 4. It still gets it right with 3.999

Hints from issue discussion:
Just a quick thought for now, though we should look at this more: it seems to be a 32bit numerical tolerance thing because enabling 64bit computation solves it:

```python
from jax import grad
from jax import config
config.update("jax_enable_x64", True)

def inv(x):
  return 1/x

f = grad(grad(grad(grad(grad(inv)))))  # one fewer `grad`!

print((f(4. + 1e-5) - f(4. - 1e-5)) / (2 * 1e-5))
print(grad(f)(4.))
```

```
0.04394531249885674
0.0439453125
```
If we look at what's going on, this appears to be a case where the higher order gradient calculation is unstable and also rather inefficient:
```
>>> make_jaxpr(grad(grad(grad(lambda x: 1/x))))(4.)
{ lambda  ; a.
  let b = mul a a
      c = mul b b
      d = div -1.0 c
      e = mul d 1.0
      f = neg e
      g = mul 1.0 f
      h = mul 1.0 f
      i = add_any g h
      j = mul 1.0 a
      k = mul 1.0 a
      l = add_any j k
      m = neg l
      n = mul m 1.0
      o = mul c c
      p = div n o
      q = mul p -1.0
      r = neg q
      s = mul r b
      t = mul r b
      u = add_any s t
      v = mul u a
      w = add_any i v
      x = mul u a
      y = add_any w x
  in (y,) }
```

The source of the issue is that `grad(inv)` turns into `-1/(x*x)` instead of `-1/x**2`:
```
>>> make_jaxpr(grad(lambda x: 1/x))(4.)
{ lambda  ; a.
  let b = mul a a
      c = div 1.0 b
      d = mul c 1.0
      e = neg d
  in (e,) }
>>> make_jaxpr(grad(lambda x: x**-1))(4.)
{ lambda  ; a.
  let b = pow a -2.0
      c = mul -1.0 b
      d = mul 1.0 c
  in (d,) }
```

The later is much more efficient for higher order differentiation:
```
>>> make_jaxpr(grad(grad(grad(lambda x: x**-1))))(4.)
{ lambda  ; a.
  let b = pow a -4.0
      c = mul -3.0 b
      d = mul 2.0 c
  in (d,) }
```

It seems like we should consider defining the gradient derivative rule for division in terms of `pow` rather than `square` (which should perhaps also be written in terms of `pow`). The `pow` primitive in turn would also probably need to be rewritten to include special case logic for efficient integer powers.
Seems this may be related to https://github.com/google/jax/pull/1388?
https://github.com/google/jax/pull/3066 is @hawkinsp's fix that adding special logic for integer scalars to `jax.numpy.power`.

Peter notes:
> I have somewhat mixed feelings about this, given its behavior isn't invariant to `jit`, but it does seem likely to improve cases like `x ** 2` if `pow` has low precision.

So I guess the ideal fix might be to upstream some of this into XLA's `pow`.
It would also be remiss not to mention `jax.experimental.jet` if you care about efficient high order differentiation. I'm sure `jet` handles this more efficiently than `grad`.
@shoyer I think we can get much the same effect by having a separate `integer_power` primitive at the lax level?
Yes, an `integer_power` primitive could work well. Possibly restricted to statically known powers?

Commit message: Add a primitive integer_pow() for values raised to fixed integer scalar.

Use integer_pow() in the RHS JVP of div(). Also use it in square() and reciprocal().

Fixes #3136

```
In [1]: from jax import grad, make_jaxpr
In [2]: def inv(x): return 1/x
In [3]: print(grad(grad(grad(grad(grad(grad(inv))))))(4.))
0.043945312

In [4]: make_jaxpr(grad(grad(grad(grad(grad(grad(inv)))))))(4.)
Out[4]:
{ lambda  ; a.
  let b = integer_pow[ y=-7 ] a
      c = mul -6.0 b
      d = mul -120.0 c
  in (d,) }

In [5]:
```

Patch to review:
diff --git a/jax/experimental/jet.py b/jax/experimental/jet.py
index d3e4d892b25b..308e00575b70 100644
--- a/jax/experimental/jet.py
+++ b/jax/experimental/jet.py
@@ -261,6 +261,15 @@ def _pow_taylor(primals_in, series_in):
   return primal_out, series_out
 jet_rules[lax.pow_p] = _pow_taylor
 
+def _integer_pow_taylor(primals_in, series_in, *, y):
+  if y == 2:
+    fn = lambda x: x * x
+  else:
+    fn = lambda x: lax.pow(x, onp.array(y, dtype=x.dtype))
+  return jet(fn, primals_in, series_in)
+jet_rules[lax.integer_pow_p] = _integer_pow_taylor
+
+
 def _expit_taylor(primals_in, series_in):
   x, = primals_in
   series, = series_in
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
index a2c3d69fedfa..5b736c4a95aa 100644
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -144,6 +144,8 @@
   index_take,
   infeed,
   infeed_p,
+  integer_pow,
+  integer_pow_p,
   iota,
   is_finite,
   is_finite_p,
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 01271dcf7dfb..f713f4120514 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -255,6 +255,15 @@ def pow(x: Array, y: Array) -> Array:
   r"""Elementwise power: :math:`x^y`."""
   return pow_p.bind(x, y)
 
+def integer_pow(x: Array, y: int) -> Array:
+  r"""Elementwise power: :math:`x^y`, where :math:`y` is a fixed integer."""
+  if y == 0:
+    return _ones(x)
+  elif y == 1:
+    return x
+  else:
+    return integer_pow_p.bind(x, y=y)
+
 def sqrt(x: Array) -> Array:
   r"""Elementwise square root: :math:`\sqrt{x}`."""
   return sqrt_p.bind(x)
@@ -1604,11 +1613,11 @@ def batch_matmul(lhs: Array, rhs: Array,
 
 def square(x: Array) -> Array:
   r"""Elementwise square: :math:`x^2`."""
-  return mul(x, x)
+  return integer_pow(x, 2)
 
 def reciprocal(x: Array) -> Array:
   r"""Elementwise reciprocal: :math:`1 \over x`."""
-  return div(_const(x, 1), x)
+  return integer_pow(x, -1)
 
 def _upcast_fp16_for_computation(f):
   @functools.wraps(f)
@@ -2054,6 +2063,41 @@ def _pow_jvp_rhs(g, ans, x, y):
   return mul(_brcast(g, x), mul(log(_replace_zero(x)), ans))
 
 ad.defjvp2(pow_p, _pow_jvp_lhs, _pow_jvp_rhs)
+
+
+def _integer_pow_dtype_rule(x, *, y):
+  dtype = unop_dtype_rule(_identity, _int | _float | _complex, 'integer_pow', x)
+  if y < 0 and dtypes.issubdtype(dtype, onp.integer):
+    raise TypeError("Integers cannot be raised to negative powers, got "
+                    f"integer_pow({x}, {y})")
+  return dtype
+
+def _integer_pow_translation_rule(c, x, *, y):
+  if y == 0:
+    shape = c.get_shape(x)
+    return xb.constant(c, onp.array(1, dtype=shape.numpy_dtype()))
+  is_reciprocal = y < 0
+  if is_reciprocal:
+    y = -y
+  acc = None
+  while y > 0:
+    if y & 1:
+      acc = x if acc is None else xops.Mul(acc, x)
+    y >>= 1
+    if y > 0:
+      x = xops.Mul(x, x)
+  return xops.Reciprocal(acc) if is_reciprocal else acc
+
+def _integer_pow_jvp(g, x, *, y):
+  return g if y == 0 else mul(g, mul(_const(x, y), integer_pow(x, y - 1)))
+
+integer_pow_p = standard_primitive(
+  _attrgetter('shape'), _integer_pow_dtype_rule, 'integer_pow',
+  translation_rule=_integer_pow_translation_rule)
+batching.defvectorized(integer_pow_p)
+masking.defvectorized(integer_pow_p)
+ad.defjvp(integer_pow_p, _integer_pow_jvp)
+
 _replace_zero = lambda x: select(eq(x, _const(x, 0)), _ones(x), x)
 
 not_p = standard_unop(_bool_or_int, 'not')
@@ -2103,7 +2147,7 @@ def _div_transpose_rule(cotangent, x, y):
 div_p = standard_naryop([_num, _num], 'div')
 ad.defjvp(div_p,
           lambda g, x, y: div(_brcast(g, y), y),
-          lambda g, x, y: div(mul(neg(_brcast(g, x)), x), square(y)))
+          lambda g, x, y: mul(mul(neg(_brcast(g, x)), x), integer_pow(y, -2)))
 ad.primitive_transposes[div_p] = _div_transpose_rule
 
 rem_p = standard_naryop([_num, _num], 'rem')
diff --git a/jax/nn/functions.py b/jax/nn/functions.py
index e9a911587d08..f37cd65578ad 100644
--- a/jax/nn/functions.py
+++ b/jax/nn/functions.py
@@ -182,9 +182,7 @@ def gelu(x):
   <https://arxiv.org/abs/1606.08415>`_, section 2.
   """
   sqrt_2_over_pi = np.sqrt(2 / np.pi).astype(x.dtype)
-  # Does not use the power operator here.
-  # See https://github.com/google/jax/pull/3036
-  x_cubed = x * x * x
+  x_cubed = lax.integer_pow(x, 3)
   cdf = 0.5 * (1.0 + jnp.tanh(sqrt_2_over_pi * (x + 0.044715 * x_cubed)))
   return x * cdf
 
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index d87dfa534fe8..5f6149e23573 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -528,16 +528,8 @@ def power(x1, x2):
   # Using lax.pow may be imprecise for floating-point values; the goal of this
   # code path is to make sure we end up with a precise output for the common
   # pattern ``x ** 2`` or similar.
-  if isinstance(x2, int) and 0 <= x2 <= 64:
-    x1 = asarray(x1)
-    acc = None
-    while x2 > 0:
-      if x2 & 1:
-        acc = x1 if acc is None else lax.mul(acc, x1)
-      x2 >>= 1
-      if x2 > 0:
-        x1 = lax.mul(x1, x1)
-    return ones_like(x1) if acc is None else acc
+  if isinstance(x2, int):
+    return lax.integer_pow(x1, x2)
 
   x1, x2 = _promote_args(np.power, x1, x2)
   dtype = _dtype(x1)
@@ -815,7 +807,7 @@ def cbrt(x):
 
 
 @_wraps(np.square)
-def square(x): return lax.mul(x, x)
+def square(x): return lax.integer_pow(x, 2)
 
 
 @_wraps(np.deg2rad)
@@ -888,7 +880,7 @@ def hypot(x1, x2):
 @_wraps(np.reciprocal)
 def reciprocal(x):
   x, = _promote_dtypes_inexact(x)
-  return lax.div(lax._const(x, 1), x)
+  return lax.integer_pow(x, -1)
 
 
 @_wraps(np.sinc, update_doc=False)
diff --git a/tests/host_callback_test.py b/tests/host_callback_test.py
index 426bab716d6f..a06c3930d826 100644
--- a/tests/host_callback_test.py
+++ b/tests/host_callback_test.py
@@ -143,7 +143,7 @@ def test_eval(self):
                     func=_print
                     nr_untapped=1
                     what=y * 3 ] d c
-      g = mul f f
+      g = integer_pow[ y=2 ] f
   in (g,) }""", str(api.make_jaxpr(fun1)(5.)))
     self.assertEqual("", testing_stream.output)
 
@@ -731,7 +731,7 @@ def test_jvp(self):
                     func=_print
                     nr_untapped=1
                     what=y * 3 ] e d
-      h = mul g g
+      h = integer_pow[ y=2 ] g
       i = mul b 2.00
       j k = id_tap[ arg_treedef=*
                     func=_print
@@ -744,10 +744,9 @@ def test_jvp(self):
                       nr_untapped=2
                       transforms=('jvp',)
                       what=y * 3 ] l j f
-      p = mul n g
-      q = mul g n
-      r = add_any p q
-  in (h, r) }""",
+      p = mul 2.00 g
+      q = mul n p
+  in (h, q) }""",
                                  str(api.make_jaxpr(jvp_fun1)(jnp.float32(5.), jnp.float32(0.1))))
     with hcb.outfeed_receiver():
       res_primals, res_tangents = jvp_fun1(jnp.float32(5.), jnp.float32(0.1))
@@ -892,7 +891,7 @@ def test_vmap(self):
                     nr_untapped=1
                     transforms=('batch',)
                     what=y * 3 ] d c
-      g = mul f f
+      g = integer_pow[ y=2 ] f
   in (g,) }""", str(api.make_jaxpr(vmap_fun1)(vargs)))
     with hcb.outfeed_receiver():
       res_vmap = vmap_fun1(vargs)
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 68600446293c..99b59259e7a0 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2668,6 +2668,14 @@ def testRemainder(self):
     tol = 1e-1 if num_float_bits(onp.float64) == 32 else 1e-3
     check_grads(lax.rem, (x, y), 2, ["fwd", "rev"], tol, tol)
 
+  def testHigherOrderGradientOfReciprocal(self):
+    # Regression test for https://github.com/google/jax/issues/3136
+    def inv(x):
+      # N.B.: intentionally written as 1/x, not x ** -1 or reciprocal(x)
+      return 1 / x
+    grad_fn = jax.grad(jax.grad(jax.grad(jax.grad(jax.grad(jax.grad(inv))))))
+    self.assertAllClose(onp.float32(0.0439453125), grad_fn(onp.float32(4.)),
+                        check_dtypes=True)
 
 def all_bdims(*shapes):
   bdims = (itertools.chain([cast(Optional[int], None)],


Code review comments:
diff_hunk: @@ -182,9 +182,7 @@ def gelu(x):
   <https://arxiv.org/abs/1606.08415>`_, section 2.
   """
   sqrt_2_over_pi = np.sqrt(2 / np.pi).astype(x.dtype)
-  # Does not use the power operator here.
-  # See https://github.com/google/jax/pull/3036
-  x_cubed = x * x * x
+  x_cubed = lax.integer_pow(x, 3)
path: jax/nn/functions.py
text: @user1:
Optional: could probably use infix `**` now:
```suggestion
  x_cubed = x ** 3
```