jax-ml/jax · error · TypeError

Called {name} with a float0 at position {i}. float0s do not

Error message

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. 
If you didn't expect to get a float0 you might have accidentally taken a gradient with respect to an integer argument.

What it means

N-ary lax ops (binops like add, mul) validate each operand position against per-position accepted dtypes. If the offending operand has dtype float0 — the tangent dtype of non-differentiable values — this specialized TypeError is raised instead of the generic dtype message, since float0 typically signals an autodiff misuse.

Source

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

                            ur_rule=partial(unop_ur_rule, name))
  batching.defvectorized(prim)
  return prim

standard_unop = partial(unop, _identity)

_attrgetter = lambda name: lambda x, **kwargs: getattr(x, name)


def naryop_dtype_rule(result_dtype, accepted_dtypes, name, *avals,
                      require_same=True, allow_extended_dtype=False, **kwargs):
  assert len(avals) == len(accepted_dtypes), (avals, accepted_dtypes)
  for i, aval in enumerate(avals):
    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. stop_gradient integer/boolean inputs so their tangents don't propagate: lax.stop_gradient(idx)
2. Cast integer primals to float before differentiating
3. Register custom VJPs that return zero tangents for non-differentiable inputs

Example fix

// before
def f(x, idx):
    return (x * w[idx]).sum()
jax.grad(f, argnums=1)(x, idx)  # float0 tangent into mul
// after
def f(x, idx):
    return (x * w[lax.stop_gradient(idx)]).sum()

import jax.numpy as jnp, jax.dtypes as dt
if any(a.dtype == dt.float0 for a in (x, y)):
    x = jnp.zeros_like(x, dtype=jnp.float32) if x.dtype == dt.float0 else x
    y = jnp.zeros_like(y, dtype=jnp.float32) if y.dtype == dt.float0 else y

def any_float0(*args) -> bool:
    dt = __import__('jax').dtypes
    return any(a.dtype == dt.float0 for a in args)

When it happens

Trigger: Using jvp/vjp/grad where a float0 tangent (from an integer or boolean primal) flows into a binary lax op at position i; chaining integer inputs through differentiable code under transformation.

Common situations: grad/jvp of functions mixing integer indices with float math; custom_vjp rules that pass raw tangents (possibly float0) into lax ops; masking pipelines where boolean masks become float0 tangents.

Related errors

• Called {name} with a float0 array. float0s do not support an
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp do not match; dtypes
• {} does not accept dtype {} at position {}. Accepted dtypes
• transpose with implicit broadcasting of unshaped values. Got

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