# Buggy code
def unop(result_dtype, accepted_dtypes, name):
    dtype_rule = partial(unop_dtype_rule, result_dtype, accepted_dtypes, name)
    prim = standard_primitive(_attrgetter("shape"), dtype_rule, name)
    batching.defvectorized(prim)
    masking.defvectorized(prim)
    return prim


# Traceback (RuntimeError)
---------------------------------------------------------------------------
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

# Fixed code
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, translation_rule=translation_rule
    )
    batching.defvectorized(prim)
    masking.defvectorized(prim)
    return prim