jax-ml/jax · error · TypeError

{} does not accept dtype {} at position {}. Accepted dtypes

Error message

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

What it means

N-ary lax ops validate each argument position's dtype against a per-position accepted_dtypes list. When operand i's dtype is not a subtype of any accepted category (and isn't float0), this TypeError reports the position and the allowed types.

Source

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

    if allow_extended_dtype and isinstance(aval.dtype, dtypes.ExtendedDType):
      continue
    types = accepted_dtypes[i]
    if not any(dtypes.issubdtype(aval.dtype, t) for t in types):
      if aval.dtype == dtypes.float0:
        raise TypeError(
            f"Called {name} with a float0 at position {i}. "
            "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.")
      else:
        msg = ('{} does not accept dtype {} at position {}. '
               'Accepted dtypes at position {} are subtypes of {}.')
        typename = dtype_to_string(aval.dtype)
        typenames = ', '.join(t.__name__ for t in types)
        raise TypeError(msg.format(name, typename, i, i, typenames))
  if require_same and kwargs.get('out_dtype') is None:
    check_same_dtypes(name, *avals)
  return result_dtype(*avals, **kwargs)


def broadcasting_shape_rule(name, *avals, **kwargs):
  if not isinstance(name, str):
    raise RuntimeError(
      "First argument of broadcasting_shape_rule should be a name."
      f" Got {name}")
  shapes = [aval.shape for aval in avals if aval.shape]
  if not shapes:
    return ()
  return _try_broadcast_shapes(*shapes, name=name)


def broadcasting_sharding_rule(name, *avals, **kwargs):
  if not isinstance(name, str):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the offending argument to an accepted dtype: x.astype(jnp.float32)
2. Check the op signature/docstring for which dtypes each position accepts
3. Ensure arrays aren't bool when arithmetic dtypes are expected (use .astype(int) first)

Example fix

// before
z = lax.integer_pow(bmask)  # bool at position 0
// after
z = lax.integer_pow(bmask.astype(jnp.int32))

x = x.astype(jnp.float32) if not jnp.issubdtype(x.dtype, jnp.inexact) else x
y = y.astype(jnp.float32) if not jnp.issubdtype(y.dtype, jnp.inexact) else y

def operands_acceptable(args, accepted) -> bool:
    import numpy as np
    return all(any(np.issubdtype(a.dtype, t) for t in ts)
               for a, ts in zip(args, accepted))

When it happens

Trigger: Passing e.g. a complex operand where only real floats are accepted at that position, or a bool where integer/float required, to lax binops like pow's operands, comparison internals, or division variants.

Common situations: Complex-valued state accidentally fed into real-only ops; bool arrays used as numbers (JAX does not implicitly promote bool to int in lax); mixed dtypes across versions where accepted lists changed.

Related errors

• {} does not accept dtype {}. Accepted dtypes are subtypes of
• 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)}.
• Called {name} with a float0 at position {i}. float0s do not

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