Repository: google/jax

Issue:
Floor division does not work with unsigned (8/16/32/64 bit) int types
```
from jax import numpy as np
np.ones((1,), np.uint32) // 2
```
Produces:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in primitive_computation(prim, *avals, **params)
    179   try:
--> 180     return c.Build()
    181   except RuntimeError as e:

8 frames
/usr/local/lib/python3.6/dist-packages/jax/lib/xla_bridge.py in Build(self, *args, **kwargs)
    256     return super(_JaxComputationBuilder, self).Build(
--> 257         *args, **kwargs)
    258 

/usr/local/lib/python3.6/dist-packages/jaxlib/xla_client.py in Build(self, root, backend)
    729     else:
--> 730       return Computation(self._builder.Build(), backend=backend)
    731 

RuntimeError: Invalid argument: Expected element type in shape to be signed or complex for sign operation; got U32.: 

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
<ipython-input-12-3db71aa0f40b> in <module>()
      1 from jax import numpy as np
      2 
----> 3 np.ones((1,), np.uint32) // 2

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in floor_divide(x1, x2)
    459   if issubdtype(dtype, integer):
    460     quotient = lax.div(x1, x2)
--> 461     select = logical_and(lax.sign(x1) != lax.sign(x2), lax.rem(x1, x2) != 0)
    462     # TODO(mattjj): investigate why subtracting a scalar was causing promotion
    463     return where(select, quotient - onp.array(1, _dtype(quotient)), quotient)

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in sign(x)
    117   \end{cases}`.
    118   """
--> 119   return sign_p.bind(x)
    120 
    121 def floor(x):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    150     top_trace = find_top_trace(args)
    151     if top_trace is None:
--> 152       return self.impl(*args, **kwargs)
    153 
    154     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    138   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    139   abstract_args = map(abstractify, args)
--> 140   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
    141   return compiled_fun(*args)
    142 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
    150   else:
    151     handle_result = aval_to_result_handler(aval_out)
--> 152   built_c = primitive_computation(prim, *abstract_args, **params)
    153   compiled = built_c.Compile(compile_options=xb.get_compile_options(),
    154                              backend=xb.get_backend(backend))

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in primitive_computation(prim, *avals, **params)
    183            "This is a bug in JAX's shape-checking rules; please report it!\n"
    184            "https://github.com/google/jax/issues\n")
--> 185     raise RuntimeError(msg)
    186 
    187 def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):

RuntimeError: Invalid argument: Expected element type in shape to be signed or complex for sign operation; got U32.: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -111,13 +111,24 @@ def neg(x):
 def sign(x):
   r"""Elementwise sign.
 
+  For floating-point inputs, returns
   :math:`\mathrm{sign}(x) = \begin{cases}
   -1 & x < 0\\
   -0 & x = -0\\
   \mathit{NaN} & x = \mathit{NaN}\\
   +0 & x = +0\\
   1 & x > 0
-  \end{cases}`.
+  \end{cases}`
+
+  For signed integer inputs, returns
+  :math:`\mathrm{sign}(x) = \begin{cases}
+  -1 & x < 0\\
+  0 & x = 0\\
+  1 & x > 0
+  \end{cases}`
+
+  For complex inputs, returns the complex phase, i.e.
+  :math:`\mathrm{sign}(x) = \frac{x}{|x|}`.
   """
   return sign_p.bind(x)
 
@@ -1531,9 +1542,10 @@ def unop_dtype_rule(result_dtype, accepted_dtypes, name, aval, **kwargs):
   return result_dtype(aval.dtype)
 
 
-def unop(result_dtype, accepted_dtypes, name):
+def unop(result_dtype, accepted_dtypes, name, translation_rule=None):
   dtype_rule = partial(unop_dtype_rule, result_dtype, accepted_dtypes, name)
-  prim = standard_primitive(_attrgetter('shape'), dtype_rule, name)
+  prim = standard_primitive(_attrgetter('shape'), dtype_rule, name,
+                            translation_rule=translation_rule)
   batching.defvectorized(prim)
   masking.defvectorized(prim)
   return prim
@@ -1623,7 +1635,17 @@ def _brcast_to(x, shape):
 neg_p = standard_unop(_num, 'neg')
 ad.deflinear(neg_p, lambda t: [neg(t)])
 
-sign_p = standard_unop(_num, 'sign')
+def _sign_translation_rule(c, x):
+  shape = c.GetShape(x)
+  dtype = shape.numpy_dtype()
+  if dtypes.issubdtype(dtype, onp.unsignedinteger):
+    zero = c.Constant(onp.array(0, dtype=dtype))
+    dims = c.GetShape(x).dimensions()
+    return c.Select(c.Eq(x, zero), c.Broadcast(zero, dims),
+                    c.Broadcast(c.Constant(onp.array(1, dtype=dtype)), dims))
+  return c.Sign(x)
+
+sign_p = standard_unop(_num, 'sign', translation_rule=_sign_translation_rule)
 ad.defjvp_zero(sign_p)
 
 nextafter_p = standard_binop(
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
@@ -415,7 +415,6 @@ def fn(x1, x2):
 bitwise_not = _one_to_one_unop(onp.bitwise_not, lax.bitwise_not)
 negative = _one_to_one_unop(onp.negative, lax.neg)
 positive = _one_to_one_unop(onp.positive, lambda x: x)
-sign = _one_to_one_unop(onp.sign, lax.sign)
 
 floor = _one_to_one_unop(onp.floor, lax.floor, True)
 ceil = _one_to_one_unop(onp.ceil, lax.ceil, True)
@@ -486,6 +485,16 @@ def op(*args):
 logical_xor = _logical_op(onp.logical_xor, lax.bitwise_xor)
 
 
+@_wraps(onp.sign)
+def sign(x):
+  dtype = _dtype(x)
+  if issubdtype(dtype, complexfloating):
+    re = lax.real(x)
+    return lax.complex(
+      lax.sign(where(re != 0, re, lax.imag(x))), _constant_like(re, 0))
+  return lax.sign(x)
+
+
 @_wraps(onp.true_divide)
 def true_divide(x1, x2):
   x1, x2 = _promote_args_inexact("true_divide", x1, x2)


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
@@ -60,6 +60,7 @@
   {lnp.bfloat16, onp.float16, onp.float32, onp.float64}))
 complex_dtypes = [onp.complex64, onp.complex128]
 int_dtypes = [onp.int32, onp.int64]
+uint_dtypes = [onp.uint32, onp.uint64]
 unsigned_dtypes = [onp.uint32, onp.uint64]
 bool_dtypes = [onp.bool_]
 default_dtypes = float_dtypes + int_dtypes
@@ -188,6 +189,10 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("fix", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("floor_divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero,
               ["rev"]),
+    # TODO(phawkins): merge this with the preceding entry after the minimum
+    # Jaxlib version is increased to 0.1.38.
+    op_record("floor_divide", 2, uint_dtypes, all_shapes, jtu.rand_nonzero,
+              ["rev"]),
     op_record("heaviside", 2, default_dtypes, all_shapes, jtu.rand_default, [],
               inexact=True),
     op_record("hypot", 2, default_dtypes, all_shapes, jtu.rand_default, [],
@@ -230,6 +235,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
               tolerance={onp.float16: 1e-2}),
     op_record("mod", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("sign", 1, number_dtypes + uint_dtypes, all_shapes,
+              jtu.rand_some_inf_and_nan, []),
     op_record("sinc", 1, [t for t in number_dtypes if t != lnp.bfloat16],
               all_shapes, jtu.rand_default, ["rev"],
               tolerance={onp.complex64: 1e-5}, inexact=True,
@@ -297,7 +304,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               tolerance={onp.float32: 2e-4, onp.complex64: 2e-4, onp.complex128: 1e-14}),
     op_record("__mod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
               tolerance={onp.float16: 1e-1}),
-    op_record("__floordiv__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("__floordiv__", 2, default_dtypes, all_shapes,
+              jtu.rand_nonzero, []),
     op_record("__truediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, [],
               inexact=True),
     op_record("__abs__", 1, number_dtypes, all_shapes, jtu.rand_default, []),
@@ -319,7 +327,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               tolerance={onp.float32: 2e-4, onp.complex64: 1e-3}),
     op_record("__rmod__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, [],
               tolerance={onp.float16: 1e-1}),
-    op_record("__rfloordiv__", 2, default_dtypes, all_shapes, jtu.rand_nonzero, []),
+    op_record("__rfloordiv__", 2, default_dtypes, all_shapes,
+              jtu.rand_nonzero, []),
     op_record("__rtruediv__", 2, number_dtypes, all_shapes, jtu.rand_nonzero, [],
               inexact=True),
     # op_record("__ror__", 2, number_dtypes, all_shapes, jtu.rand_bool, []),
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -62,6 +62,7 @@ def num_float_bits(dtype):
     {onp.complex64, onp.complex128}))
 inexact_dtypes = float_dtypes + complex_dtypes
 int_dtypes = list(jtu.supported_dtypes().intersection({onp.int32, onp.int64}))
+uint_dtypes = list(jtu.supported_dtypes().intersection({onp.uint32, onp.uint64}))
 bool_dtypes = [onp.bool_]
 default_dtypes = float_dtypes + int_dtypes
 all_dtypes = float_dtypes + complex_dtypes + int_dtypes + bool_dtypes
@@ -77,7 +78,7 @@ def op_record(op, nargs, dtypes, rng_factory, tol=None):
 
 LAX_OPS = [
     op_record("neg", 1, default_dtypes + complex_dtypes, jtu.rand_small),
-    op_record("sign", 1, default_dtypes, jtu.rand_small),
+    op_record("sign", 1, default_dtypes + uint_dtypes, jtu.rand_small),
     op_record("floor", 1, float_dtypes, jtu.rand_small),
     op_record("ceil", 1, float_dtypes, jtu.rand_small),
     op_record("round", 1, float_dtypes, jtu.rand_default),