jax-ml/jax · error · TypeError

{} does not accept dtype {}. Accepted dtypes are subtypes of

Error message

{} does not accept dtype {}. Accepted dtypes are subtypes of {}.

What it means

JAX lax elementwise ops (unops built via unop_dtype_rule, e.g. exp, floor, conj) restrict their input dtypes to accepted_dtypes (typically subtypes of floating/complex/integer). If the input dtype is not a subtype of any accepted category, this TypeError is raised listing what is allowed.

Source

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

_strip_weak_type = lambda *args, **_: False


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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the input to an accepted dtype first: x.astype(jnp.float32)
2. Check the op's docstring for accepted dtypes and pick a supported variant (e.g. use float input for exp/log)
3. For bool inputs, convert with .astype(jnp.float32) before math ops

Example fix

// before
y = lax.exp(int_array)
// after
y = lax.exp(int_array.astype(jnp.float32))

import jax.numpy as jnp
if not jnp.issubdtype(x.dtype, jnp.floating):
    x = x.astype(jnp.float32)
out = lax.exp(x)

def has_dtype_subtype(x, categories) -> bool:
    import numpy as np
    return any(np.issubdtype(x.dtype, t) for t in categories)

When it happens

Trigger: Passing an unsupported dtype to a lax unop: e.g. a string/object dtype, bool where only inexact accepted, or a truncated float16 where the op's accepted list excludes it.

Common situations: Feeding object/string arrays from data loading into lax math; applying float-only ops (expm1, tanh on some paths) to integers when the op requires inexact; using custom dtypes not registered with the op.

Related errors

• {} does not accept dtype {} at position {}. Accepted dtypes
• full must be called with scalar fill_value, got fill_value.s
• offset must be an integer, got {offset!r}
• {name} does not accept dtype {dtype_to_string(aval.dtype)}.
• the first argument to pow must have an inexact dtype (float

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