jax-ml/jax · error · TypeError

Called {name} with a float0 array. float0s do not support an

Error message

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

What it means

float0 is JAX's dtype for tangent/gradient values of non-differentiable (integer/boolean) inputs. By design float0 supports no operations; passing such an array into a lax unop (e.g. exp, neg) raises this TypeError rather than doing any implicit conversion.

Source

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

    dtypes.python_scalar_types, array_types, [array.ArrayImpl],
    literals.typed_scalar_types):
  ad_util.raw_jaxval_adders[t] = _add_arrays


### primitives


_fixed_dtype = \
    lambda dtype: lambda *args, **kwargs: np.dtype(dtype)
_complex_basetype = lambda dtype, **kwargs: np.abs(np.zeros((), dtype)).dtype

_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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Differentiate only with respect to float arguments; make the arg a float (e.g. 1.0 * x) if its gradient is meaningful
2. In custom rules, stop gradients on integer paths with lax.stop_gradient
3. If you genuinely want zeros, convert: jnp.zeros_like(arr, dtype=jnp.float32) as the message suggests

Example fix

// before
g = jax.grad(lambda i: f(i).sum())(n_ints)  # tangents are float0
// after
g = jax.grad(lambda x: f(x).sum())(n_ints.astype(jnp.float32))

import jax.numpy as jnp, jax.dtypes as dt
if x.dtype == dt.float0:
    x = jnp.zeros_like(x, dtype=jnp.float32)

def is_float0(x) -> bool:
    return x.dtype == __import__('jax').dtypes.float0

When it happens

Trigger: Differentiating (grad/jacfwd/jvp) a function with respect to an integer or boolean argument and then applying a lax elementwise op to the tangent; forwarding float0 cotangents into arithmetic inside custom VJP rules.

Common situations: Calling jax.grad on a function whose argument is an index array, count, or mask; grad through integer state updates; custom_vjp rules forgetting to zero out integer-arg tangents.

Related errors

• Called {name} with a float0 at position {i}. float0s do not
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp do not match; dtypes
• Buffer callbacks do not support transpose. Please use `jax.c
• transpose with implicit broadcasting of unshaped values. Got

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