Fix JVP rule for lax.pow()

jax/_src/lax/lax.py

@@ -2001,8 +2001,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))

tests/lax_autodiff_test.py

@@ -529,6 +529,17 @@ 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))
+
+ result_fwd = jax.jacfwd(jax.jacfwd(lax.pow, (0, 1)), (0, 1))(x, y)
+ self.assertAllClose(result_fwd, expected)
+
+ result_rev = jax.jacrev(jax.jacrev(lax.pow, (0, 1)), (0, 1))(x, y)
+ self.assertAllClose(result_rev, expected)
+
  @parameterized.named_parameters(jtu.cases_from_list(
  {"testcase_name": "_predshape={}_argshapes={}".format(
  jtu.format_shape_dtype_string(pred_shape, np.bool_),