Repository: google/jax

Issue:
pow grad has surprising behavior based on dtype
> Hi @mattjj, thanks for this. I'm finding this choice very confusing, since at the back of my head, I'm dealing with a function in one variable, x, and so it's unexpected that for every expression I write down inside such a function, I need to check to make sure what I think are constants are well-behaved if they were to become variables. But perhaps there are good reasons why this should be. 
> 
> I'm still running into trouble along these lines though, I think when `arange` is used to build the polynomial. Here is a simple example that breaks: 
```
import jax
import jax.numpy as jnp

b = lambda z: jnp.sum(z**jnp.arange(0, 2))
d = lambda z: 1. + z
print(b(2.543), d(2.543)) # check they are the same and I'm not crazy
print(jax.grad(d)(1.))  # gives 1. as expected
print(jax.grad(b)(1.)) # gives 2.
```

_Originally posted by @hongwanliu in https://github.com/google/jax/issues/14397#issuecomment-1703838855_
            


Discussion:
Copying a comment from #14397 

> If you change the `jnp.arange` to use float dtype, then things agree:
> 
> ```python
> b = lambda z: jnp.sum(z**jnp.arange(0, 2, dtype=float))
> d = lambda z: 1 + z
> 
> print(jax.grad(d)(1.))  # gives 1. as expected
> print(jax.grad(b)(1.)) # gives 1. also
> ```
> 
> I'd say that's pretty surprising though. We probably need to revise this; either there's some bug, or this dtype-based-resolving-of-ambiguities is too subtle and we should raise an error instead (asking the user to be explicit).


Wow, this was just a bug in the JVP rule, and the test I added had the same bug. Sorry! This is what I get for context switching on that PR so many times.

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
@@ -1993,7 +1993,7 @@ def _pow_jvp_lhs(g, ans, x, y):
       shape = broadcast_shapes(x.shape, y.shape)
       x = _maybe_broadcast(shape, x)
       y = _maybe_broadcast(shape, y)
-    jac = select(eq(y, _const(y, 0)), _ones(y),
+    jac = select(eq(y, _const(y, 0)), _zeros(y),
                  mul(_replace_zero(y), pow(x, sub(y, _ones(y)))))
   else:
     jac = mul(y, pow(x, sub(y, _ones(y))))


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
@@ -539,7 +539,22 @@ def testPowSecondDerivative(self):
   def testPowIntPowerAtZero(self):
     # https://github.com/google/jax/issues/14397
     ans = jax.grad(jax.jit(lambda x, n: x ** n))(0., 0)
-    self.assertAllClose(ans, 1., check_dtypes=False)
+    self.assertAllClose(ans, 0., check_dtypes=False)
+
+  @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises mixed type promotion
+  def testPowIntPowerAtZero2(self):
+    # https://github.com/google/jax/issues/17995
+    a = lambda z: jax.numpy.sum(z**jax.numpy.arange(0, 2, dtype=int))
+    b = lambda z: jax.numpy.sum(z**jax.numpy.arange(0, 2, dtype=float))
+    c = lambda z: 1 + z
+    d = lambda z: z ** 0 + z
+    e = lambda z: z ** 0. + z
+
+    self.assertAllClose(jax.grad(a)(3.14), 1., check_dtypes=False)
+    self.assertAllClose(jax.grad(b)(3.14), 1., check_dtypes=False)
+    self.assertAllClose(jax.grad(c)(3.14), 1., check_dtypes=False)
+    self.assertAllClose(jax.grad(d)(3.14), 1., check_dtypes=False)
+    self.assertAllClose(jax.grad(e)(3.14), 1., check_dtypes=False)
 
   @jtu.sample_product(
     [dict(arg_shape=arg_shape, pred_shape=pred_shape)