jnp.abs(int4) raises MLIR error — should we improve error messages for limited dtypes like int4/uint4?

[Kairoven]

Context

I recently tried using jnp.int4 and jnp.uint4 in a computation:

import jax.numpy as jnp

x = jnp.array([1, 2, 2], dtype=jnp.int4)
print(jnp.abs(x))  # ❌ Fails with MLIR verification error

This raises a very low-level error:

MLIRError: Verification failed:
error: type of return operand 0 ('tensor<3xi4>') doesn't match function result type ('tensor<3xi8>') in function @abs
...
ValueError: Cannot lower jaxpr with verifier errors

On the other hand:

x = jnp.array([1, 2, 2], dtype=jnp.uint4)
print(jnp.abs(x))  # ✅ Works and returns [1 2 2]

I understand from previous discussions that int4 and uint4 are experimental and not fully supported across JAX — that’s fine.
But the error message here is problematic. Should JAX provide clearer feedback when operations on experimental types like int4/uint4 are not supported?

[jakevdp]

Thanks for raising this issue: we should be able to catch errors like this during abstract evaluation. I've created #31644 from this discussion to track work on this.

[Kairoven]

much appreciated!