Repository: google/jax

Issue:
Exponentiation on int32 data with base 0 gives different results when jitted
Example

```python
from jax import numpy as jnp
import jax

EXP_MULTIPLIER = 1


def power(a, b):
    return a**b


a = jnp.int32(0)
b = jnp.int32(EXP_MULTIPLIER * 64)

result = power(a, b)
result_jit = jax.jit(power)(a, b)

assert result == result_jit, f'{result} != {result_jit}'
```

```
AssertionError: 0 != 1
```


Discussion:
That's very curious - thanks for the report
```python
from jax import jit
import jax.numpy as jnp

print(jnp.power(0, 64))
# 0
print(jit(jnp.power)(0, 64))
# 1
```
That we observe differences between `jit` and non-`jit` isn't surprising: that's the difference between `lax.integer_pow` and the implementation in `lax_numpy` for specific unknown integers. But why is the latter implementation doing that?
Ah: it only looks at the first 6 bits of the exponent, but in this special case (0) we actually need to perform at least one multiplication.
Wow this is nasty :)

@hawkinsp, @jakevdp do you have an approximate timeline for when this might be fixed?
`x ** y` overflows for any integer `x > 1` and `y >= 64` and therefore has undefined output under those conditions. We are just missing a special case for `x == 0`. I think the fix is trivial and it is just to add `lax.select(x == 0, 0, ...)` to the return value.

 Jake self-assigned this though so I'll let him look at it.
I'll try to get a fix in this morning.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -607,12 +607,14 @@ def power(x1, x2):
 
   # TODO(phawkins): add integer pow support to XLA.
   bits = 6  # Anything more would overflow for any x1 > 1
-  acc = ones(shape(x1), dtype=dtype)
+  zero = _constant_like(x2, 0)
+  one = _constant_like(x2, 1)
+  # Initialize acc carefully such that pow(0, x2) is zero for x2 != 0
+  acc = where(lax.bitwise_and(lax.eq(x1, zero), lax.ne(x2, zero)), zero, one)
   for _ in range(bits):
-    acc = where(lax.bitwise_and(x2, _constant_like(x2, 1)),
-                lax.mul(acc, x1), acc)
+    acc = where(lax.bitwise_and(x2, one), lax.mul(acc, x1), acc)
     x1 = lax.mul(x1, x1)
-    x2 = lax.shift_right_logical(x2, _constant_like(x2, 1))
+    x2 = lax.shift_right_logical(x2, one)
   return acc
 
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1754,6 +1754,16 @@ def testIntegerPower(self, ptype):
     self.assertLen(eqns, 1)
     self.assertEqual(eqns[0].primitive, lax.integer_pow_p)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_x={}_y={}".format(x, y), "x": x, "y": y}
+      for x in [-1, 0, 1]
+      for y in [0, 32, 64, 128]))
+  def testIntegerPowerOverflow(self, x, y):
+    # Regression test for https://github.com/google/jax/issues/5987
+    args_maker = lambda: [x, y]
+    self._CheckAgainstNumpy(np.power, jnp.power, args_maker)
+    self._CompileAndCheck(jnp.power, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis),