Repository: google/jax

Issue:
The second order gradient of `lax.pow` is wrong
The second order gradient of `lax.pow` is wrong. For instance, the `fn` computes `x^y`, so `grad_fn` returns `(y*x^(y-1), ln(x)* x^y)`. When `x, y = 23.0, 0.0`, for the second element `ln(x) * x^y`, the gradient w.r.t `x` should be `1/x` instead of `23.0` computed by `jacrev` or `jacfwd`. 

```python
import jax
def grad_fn(x, y):
    def fn(x, y):
        return jax.numpy.sum(jax.lax.pow(x, y))
    return jax.grad(fn, (0, 1))(x, y)
x = jax.numpy.array([23.0])
y = jax.numpy.array([0.0])

print(jax.jacrev(grad_fn, (0, 1))(x, y))
print(jax.jacfwd(grad_fn, (0, 1))(x, y))
```

```
((DeviceArray([[0.]], dtype=float32), DeviceArray([[23.]], dtype=float32)), (DeviceArray([[0.04347826]], dtype=float32), DeviceArray([[9.831324]], dtype=float32)))
```


Discussion:
Thanks for the report – I had a bit of trouble following your reproduction, so I rewrote it a bit more succinctly:
```python
import jax

f = jax.lax.pow
x = 23.0
y = 0.0

print(f"d²f/dxdy = {jax.grad(jax.grad(f, 0), 1)(x, y)}")
print(f"d²f/dydx = {jax.grad(jax.grad(f, 1), 0)(x, y)}")
```
```
d²f/dxdy = 23.0
d²f/dydx = 0.043478261679410934
```
As you mention, these two results should both be equal to the second output.

Looking into the code, I think the issue comes from the `select` statement in line 2004 here: https://github.com/google/jax/blob/da4e79a6251ae71504061b496ea34f20dc06fa60/jax/_src/lax/lax.py#L2003-L2008

The gradient of `pow` with respect to `x` is special-cased when `y` is zero. I think this yields the correct result when `y` is a constant, but it removes the dependence of the result on `y` and therefore the mixed second derivative is incorrect.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2004,8 +2004,7 @@ def _abs_jvp_rule(g, ans, x):
 pow_p = standard_naryop([_float | _complex, _float | _complex], 'pow')
 
 def _pow_jvp_lhs(g, ans, x, y):
-  jac = mul(y, pow(x, select(eq(y, _zeros(y)), _ones(y), sub(y, _ones(y)))))
-  return mul(g, jac)
+  return mul(g, mul(y, pow(x, sub(y, _ones(y)))))
 
 def _pow_jvp_rhs(g, ans, x, y):
   return mul(g, mul(log(_replace_zero(x)), ans))


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -529,6 +529,26 @@ def testReverseGrad(self):
     check_grads(rev, (np.array([[6., 5., 4.], [3., 2., 1.]]),), 2,
                 rtol={np.float32: 3e-3})
 
+  def testPowSecondDerivative(self):
+    # https://github.com/google/jax/issues/12033
+    x, y = 4.0, 0.0
+    expected = ((0.0, 1/x), (1/x, np.log(x) ** 2))
+
+    with self.subTest("jacfwd"):
+      result_fwd = jax.jacfwd(jax.jacfwd(lax.pow, (0, 1)), (0, 1))(x, y)
+      self.assertAllClose(result_fwd, expected)
+
+    with self.subTest("jacrev"):
+      result_rev = jax.jacrev(jax.jacrev(lax.pow, (0, 1)), (0, 1))(x, y)
+      self.assertAllClose(result_rev, expected)
+
+    with self.subTest("zero to the zero"):
+      result = jax.grad(lax.pow)(0.0, 0.0)
+      # TODO(jakevdp) special-case zero in a way that doesn't break other cases
+      # See https://github.com/google/jax/pull/12041#issuecomment-1222766191
+      # self.assertEqual(result, 0.0)
+      self.assertAllClose(result, np.nan)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_predshape={}_argshapes={}".format(
           jtu.format_shape_dtype_string(pred_shape, np.bool_),