Repository: google/jax

Issue:
jnp.right_shift incorrect on unsigned ints
It looks like ~~`lax.shift_right_arithmetic`~~ `jnp.right_shift` incorrectly treats unsigned ints as signed:

```python
import numpy as np
from jax import lax, numpy as jnp
args = np.uint8(0b10000000), np.uint8(2)
print(f"{np.right_shift(*args):#010b} (expected)")
print(f"{jnp.right_shift(*args):#010b}")
print(f"{lax.shift_right_arithmetic(*args):#010b}")
```

results in

```
0b00100000 (expected)
0b11100000
0b11100000
```

Other unsigned int types produce the same issue. Tests for `lax.shift_right_arithmetic` were missing, they will be added with https://github.com/google/jax/pull/3923.


Discussion:
I think that's working as intended. `shift_right_arithmetic` does an arithmetic (signed) shift irrespective of the type. If you want a logical shift you should use `shift_right_logical`. `lax` follows XLA in making the type of the shift independent of the carrier type.
@hawkinsp In that case, shouldn't `jnp.right_shift` resolve to `lax.shift_right_logical`  for unsigned ints to match the behavior of `np`? 

I could write the [missing `jnp.right_shift`/`left_shift` tests](https://github.com/google/jax/blob/master/tests/lax_numpy_test.py#L354) and fix this.
Yes, the jnp behavior should match the np behavior when possible. A PR with a fix & tests would be a great contribution!

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -506,7 +506,9 @@ def __xor__(self, other): return self.aval._xor(self, other)
   def __rxor__(self, other): return self.aval._rxor(self, other)
   def __invert__(self): return self.aval._invert(self)
   def __lshift__(self, other): return self.aval._lshift(self, other)
+  def __rlshift__(self, other): return self.aval._rlshift(self, other)
   def __rshift__(self, other): return self.aval._rshift(self, other)
+  def __rrshift__(self, other): return self.aval._rrshift(self, other)
   def __getitem__(self, idx): return self.aval._getitem(self, idx)
   def __nonzero__(self): return self.aval._nonzero(self)
   def __bool__(self): return self.aval._bool(self)
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -394,7 +394,6 @@ def fn(x1, x2):
 bitwise_and = _one_to_one_binop(np.bitwise_and, lax.bitwise_and)
 bitwise_or = _one_to_one_binop(np.bitwise_or, lax.bitwise_or)
 bitwise_xor = _one_to_one_binop(np.bitwise_xor, lax.bitwise_xor)
-right_shift = _one_to_one_binop(np.right_shift, lax.shift_right_arithmetic)
 left_shift = _one_to_one_binop(np.left_shift, lax.shift_left)
 equal = _one_to_one_binop(np.equal, lax.eq)
 multiply = _maybe_bool_binop(np.multiply, lax.mul, lax.bitwise_and)
@@ -441,6 +440,14 @@ def op(*args):
 logical_xor = _logical_op(np.logical_xor, lax.bitwise_xor)
 
 
+@_wraps(np.right_shift)
+def right_shift(x1, x2):
+  x1, x2 = _promote_args(np.right_shift.__name__, x1, x2)
+  lax_fn = lax.shift_right_logical if \
+    np.issubdtype(x1.dtype, np.unsignedinteger) else lax.shift_right_arithmetic
+  return lax_fn(x1, x2)
+
+
 @_wraps(np.absolute)
 def absolute(x):
   return x if issubdtype(_dtype(x), unsignedinteger) else lax.abs(x)
@@ -4492,6 +4499,8 @@ def _operator_round(number, ndigits=None):
     "invert": bitwise_not,
     "lshift": _defer_to_unrecognized_arg(left_shift),
     "rshift": _defer_to_unrecognized_arg(right_shift),
+    "rlshift": _defer_to_unrecognized_arg(_swap_args(left_shift)),
+    "rrshift": _defer_to_unrecognized_arg(_swap_args(right_shift)),
     "round": _operator_round,
 }
 


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
@@ -69,6 +69,9 @@
 
 python_scalar_dtypes = [jnp.bool_, jnp.int_, jnp.float_, jnp.complex_]
 
+# uint64 is problematic because with any uint type it promotes to float:
+int_dtypes_no_uint64 = [d for d in int_dtypes + unsigned_dtypes if d != np.uint64]
+
 def _valid_dtypes_for_shape(shape, dtypes):
   # Not all (shape, dtype) pairs are valid. In particular, Python scalars only
   # have one type in each category (float, bool, etc.)
@@ -351,7 +354,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     # op_record("__and__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
     # op_record("__xor__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
     # op_record("__divmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, []),
-    # TODO(mattjj): lshift, rshift
+    op_record("__lshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
+    op_record("__rshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
 ]
 
 JAX_RIGHT_OPERATOR_OVERLOADS = [
@@ -370,6 +374,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     # op_record("__rand__", 2, number_dtypes, all_shapes, jtu.rand_default, []),
     # op_record("__rxor__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
     # op_record("__rdivmod__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("__rlshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), []),
+    op_record("__rrshift__", 2, int_dtypes_no_uint64, all_shapes, partial(jtu.rand_int, high=8), [])
 ]
 
 class _OverrideEverything(object):
@@ -392,11 +398,8 @@ class _OverrideNothing(object):
   JAX_COMPOUND_OP_RECORDS += [
       op_record("isclose", 2, [t for t in all_dtypes if t != jnp.bfloat16],
                 all_shapes, jtu.rand_small_positive, []),
-      # uint64 is problematic because with any other int type it promotes to float.
-      op_record("gcd", 2, [d for d in int_dtypes + unsigned_dtypes if d != np.uint64],
-                all_shapes, jtu.rand_default, []),
-      op_record("lcm", 2, [d for d in int_dtypes + unsigned_dtypes if d != np.uint64],
-                all_shapes, jtu.rand_default, []),
+      op_record("gcd", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
+      op_record("lcm", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
   ]
   JAX_REDUCER_NO_DTYPE_RECORDS += [
       op_record("ptp", 1, number_dtypes, nonempty_shapes, jtu.rand_default, []),
@@ -594,6 +597,35 @@ def testBitwiseOp(self, np_op, jnp_op, rng_factory, shapes, dtypes):
                             check_dtypes=jtu.PYTHON_SCALAR_SHAPE not in shapes)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": jtu.format_test_name_suffix(op.__name__, shapes, dtypes),
+     "op": op, "dtypes": dtypes, "shapes": shapes}
+    for op in [jnp.left_shift, jnp.right_shift]
+    for shapes in filter(
+      _shapes_are_broadcast_compatible,
+      # TODO numpy always promotes to shift dtype for zero-dim shapes:
+      itertools.combinations_with_replacement(nonzerodim_shapes, 2))
+    for dtypes in itertools.product(
+      *(_valid_dtypes_for_shape(s, int_dtypes_no_uint64) for s in shapes))))
+  def testShiftOpAgainstNumpy(self, op, dtypes, shapes):
+    dtype, shift_dtype = dtypes
+    signed_mix = np.issubdtype(dtype, np.signedinteger) != \
+                 np.issubdtype(shift_dtype, np.signedinteger)
+    has_32 = any(np.iinfo(d).bits == 32 for d in dtypes)
+    promoting_to_64 = has_32 and signed_mix
+    if promoting_to_64 and not FLAGS.jax_enable_x64:
+      self.skipTest("np.right_shift/left_shift promoting to int64"
+                    "differs from jnp in 32 bit mode.")
+
+    info, shift_info = map(np.iinfo, dtypes)
+    x_rng = jtu.rand_int(self.rng(), low=info.min, high=info.max + 1)
+    # NumPy requires shifts to be non-negative and below the bit width:
+    shift_rng = jtu.rand_int(self.rng(), high=max(info.bits, shift_info.bits))
+    args_maker = lambda: (x_rng(shapes[0], dtype), shift_rng(shapes[1], shift_dtype))
+    self._CompileAndCheck(op, args_maker)
+    np_op = getattr(np, op.__name__)
+    self._CheckAgainstNumpy(np_op, op, args_maker)
+
   @parameterized.named_parameters(itertools.chain.from_iterable(
       jtu.cases_from_list(
         {"testcase_name": "{}_inshape={}_axis={}_dtype={}_keepdims={}".format(