Repository: google/jax

Issue:
lax.abs crashes on unsigned ints
Repro:
```python
import jax
jax.lax.abs(jax.numpy.uint32(1))
```
```pytb
Traceback (most recent call last):
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/mlir.py", line 814, in lower_jaxpr_to_module
    if not ctx.module.operation.verify():
jaxlib.mlir._mlir_libs.MLIRError: Verification failed:
error: "jit(abs)/jit(main)/abs"("/Users/vanderplas/github/google/jax/tmp.py":2:0): 'stablehlo.abs' op operand #0 must be tensor of 4/8/16/32/64-bit signless integer or f8E4M3B11FNUZ type or f8E4M3FN type or f8E4M3FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 4/8/16/32-bit uniform quantized signed integer or 4/8/16/32-bit uniform quantized unsigned integer values, but got 'tensor<ui32>'
 note: "jit(abs)/jit(main)/abs"("/Users/vanderplas/github/google/jax/tmp.py":2:0): see current operation: %0 = "stablehlo.abs"(%arg0) : (tensor<ui32>) -> tensor<ui32>

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/vanderplas/github/google/jax/tmp.py", line 2, in <module>
    jax.lax.abs(jax.numpy.uint32(1))
  File "/Users/vanderplas/github/google/jax/jax/_src/lax/lax.py", line 368, in abs
    return abs_p.bind(x)
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 386, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 389, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 869, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/Users/vanderplas/github/google/jax/jax/_src/dispatch.py", line 128, in apply_primitive
    compiled_fun = xla_primitive_callable(
  File "/Users/vanderplas/github/google/jax/jax/_src/util.py", line 263, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/_src/util.py", line 256, in cached
    return f(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/_src/dispatch.py", line 157, in xla_primitive_callable
    computation = sharded_lowering(
  File "/Users/vanderplas/github/google/jax/jax/_src/dispatch.py", line 188, in sharded_lowering
    return pxla.lower_sharding_computation(
  File "/Users/vanderplas/github/google/jax/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/pxla.py", line 2049, in lower_sharding_computation
    nreps, tuple_args, shape_poly_state) = _cached_lowering_to_hlo(
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/pxla.py", line 1850, in _cached_lowering_to_hlo
    lowering_result = mlir.lower_jaxpr_to_module(
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/mlir.py", line 829, in lower_jaxpr_to_module
    raise ValueError("\n".join(msg_lines)) from e
ValueError: Cannot lower jaxpr with verifier errors:
	'stablehlo.abs' op operand #0 must be tensor of 4/8/16/32/64-bit signless integer or f8E4M3B11FNUZ type or f8E4M3FN type or f8E4M3FNUZ type or f8E5M2 type or f8E5M2FNUZ type or 16-bit float or 32-bit float or 64-bit float or bfloat16 type or complex type with 32-bit float or 64-bit float elements or 4/8/16/32-bit uniform quantized signed integer or 4/8/16/32-bit uniform quantized unsigned integer values, but got 'tensor<ui32>'
		at loc("jit(abs)/jit(main)/abs"("/Users/vanderplas/github/google/jax/tmp.py":2:0))
	see current operation: %0 = "stablehlo.abs"(%arg0) : (tensor<ui32>) -> tensor<ui32>
		at loc("jit(abs)/jit(main)/abs"("/Users/vanderplas/github/google/jax/tmp.py":2:0))
Module string:
#loc = loc(unknown)
"builtin.module"() <{sym_name = "jit_abs"}> ({
  "func.func"() <{arg_attrs = [{mhlo.sharding = "{replicated}"}], function_type = (tensor<ui32>) -> tensor<ui32>, res_attrs = [{}], sym_name = "main", sym_visibility = "public"}> ({
  ^bb0(%arg0: tensor<ui32> loc(unknown)):
    %0 = "stablehlo.abs"(%arg0) : (tensor<ui32>) -> tensor<ui32> loc(#loc2)
    "func.return"(%0) : (tensor<ui32>) -> () loc(#loc)
  }) : () -> () loc(#loc)
}) {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} : () -> () loc(#loc)
#loc1 = loc("/Users/vanderplas/github/google/jax/tmp.py":2:0)
#loc2 = loc("jit(abs)/jit(main)/abs"(#loc1))
```
I think the fix is to avoid binding the primitive for unsigned inputs.


Discussion:

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
@@ -1710,6 +1710,7 @@ def _nary_lower_hlo(op: Callable, ctx,
 _complex_elem_types = {np.float32, np.float64}
 _int = {np.integer}
 _bool = {np.bool_}
+_signedint = {np.signedinteger}
 
 _num = _int | _float | _complex
 _any = _int | _float | _complex | _bool
@@ -1944,7 +1945,7 @@ def _conj_transpose_rule(t, x, *, input_dtype):
 ad.primitive_jvps[conj_p] = partial(ad.linear_jvp, conj_p)
 ad.primitive_transposes[conj_p] = _conj_transpose_rule
 
-abs_p = unop(_complex_basetype, _num, 'abs')
+abs_p = unop(_complex_basetype, _signedint | _float | _complex, 'abs')
 mlir.register_lowering(abs_p, partial(_nary_lower_hlo, hlo.AbsOp))
 
 def _abs_jvp_rule(g, ans, x):


Test patch: