jax-ml/jax · error · TypeError

{name} does not accept dtype {dtype_to_string(aval.dtype)}.

Error message

{name} does not accept dtype {dtype_to_string(aval.dtype)}. Support for narrow-width integers is platform-dependent and limited to a few specific operations, e.g. basic arithmetic and type casting.

What it means

Some lax ops disallow narrow-width integer dtypes (uint2/int2/uint4/int4) because backend support is platform-dependent. When supports_narrow_ints=False and the input has one of these dtypes, this TypeError is raised.

Source

Thrown at jax/_src/lax/lax.py:4264

def unop_dtype_rule(result_dtype, accepted_dtypes, name, aval,
                    supports_narrow_ints=True, **kwargs):
  if aval.dtype == dtypes.float0:
    raise TypeError(
        f"Called {name} with a float0 array. "
        "float0s do not support any operations by design, because they "
        "are not compatible with non-trivial vector spaces. No implicit dtype "
        "conversion is done. You can use np.zeros_like(arr, dtype=np.float) "
        "to cast a float0 array to a regular zeros array. \n"
        "If you didn't expect to get a float0 you might have accidentally "
        "taken a gradient with respect to an integer argument.")
  if not any(dtypes.issubdtype(aval.dtype, t) for t in accepted_dtypes):
    msg = '{} does not accept dtype {}. Accepted dtypes are subtypes of {}.'
    typename = dtype_to_string(aval.dtype)
    accepted_typenames = (t.__name__ for t in accepted_dtypes)
    raise TypeError(msg.format(name, typename, ', '.join(accepted_typenames)))
  if (not supports_narrow_ints) and aval.dtype in [dtypes.uint2, dtypes.int2, dtypes.uint4, dtypes.int4]:
    raise TypeError(f'{name} does not accept dtype {dtype_to_string(aval.dtype)}.'
                    ' Support for narrow-width integers is platform-dependent'
                    ' and limited to a few specific operations, e.g. basic'
                    ' arithmetic and type casting.')
  return result_dtype(aval.dtype, **kwargs)

def default_unop_reduced_rule(aval):
  return getr(aval)

def unop_ur_rule(name, aval, **kwargs):
  reduced = default_unop_reduced_rule(aval)
  if any(getu(aval)):
    raise NotImplementedError(
        f'unreduced rule for {name} is not implemented. Please'
        ' file an issue at https://github.com/jax-ml/jax/issues')
  return frozenset(), reduced, None

def unop(result_dtype, accepted_dtypes, name, supports_narrow_ints=True):
  dtype_rule = partial(unop_dtype_rule, result_dtype, accepted_dtypes, name,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Widen the values before the op: x.astype(jnp.int32) or jnp.float32, then narrow back if needed
2. Use only the supported ops (basic arithmetic and casting) on narrow dtypes, per the message
3. Keep narrow dtypes confined to storage; compute in int8/int32

Example fix

// before
y = lax.some_op(x_int4)
// after
y = lax.some_op(x_int4.astype(jnp.int32)).astype(jnp.int4)

NARROW = {jnp.uint2, jnp.int2, jnp.uint4, jnp.int4}
if x.dtype in NARROW:
    x = x.astype(jnp.int32)
out = lax.some_op(x)

def is_narrow_int(x) -> bool:
    import jax.numpy as jnp
    return x.dtype in {jnp.uint2, jnp.int2, jnp.uint4, jnp.int4}

When it happens

Trigger: Passing int2/uint2/int4/uint4 arrays into lax ops that do not opt into narrow-int support (most transcendental/unusual unops).

Common situations: Using experimental 4-bit/2-bit quantized activations with the full lax op set; importing models that stored weights in sub-byte dtypes and calling arbitrary lax ops on them.

Related errors

• {} does not accept dtype {}. Accepted dtypes are subtypes of
• {} does not accept dtype {} at position {}. Accepted dtypes
• lax.platform_dependent: the '{pname}' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• lax.platform_dependent: the 'default' branch must be a calla

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8f37973b14b91d0d. Report an issue: GitHub.