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 number.

What it means

The cumulative-reduction dtype rule only accepts numeric dtypes (subtypes of numpy.number). Non-numeric operands — bool is the classic case for cumsum/cumlogsumexp in some versions, or strings/objects — fail this check and raise TypeError naming the primitive and dtype.

Source

Thrown at jax/_src/lax/control_flow/loops.py:3094

    return x
  padding = [(0, 0)] * x.ndim
  padding[axis] = (0, n - 1) if reverse else (n - 1, 0)
  strides = [1] * x.ndim
  window_dims = [1] * x.ndim
  window_dims[axis] = n
  return window_reduce(x, window_dims, strides, padding)


def _cumred_batch_rule(prim, batched_args, batch_dims, *, axis: int,
                       reverse: bool):
  operand, = batched_args
  bdim, = batch_dims
  axis = axis if axis < bdim else axis + 1
  return prim.bind(operand, axis=axis, reverse=reverse), bdim

def _cumred_dtype_rule(name, operand, *args, **kw):
  if not dtypes.issubdtype(operand.dtype, np.number):
    raise TypeError("{} does not accept dtype {}. Accepted dtypes are subtypes "
                    "of number.".format(name, np.dtype(operand.dtype).name))
  return operand.dtype


def _cumulative_reduction_primitive(name, reduce_fn, reduce_window_fn):
  reducer_p = lax.standard_primitive(
    _cumred_shape_rule, partial(_cumred_dtype_rule, name),
    name, sharding_rule=_cumred_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, name))
  batching.primitive_batchers[reducer_p] = partial(_cumred_batch_rule,
                                                   reducer_p)

  def register_lowering(fn, platform=None):
    mlir.register_lowering(
        reducer_p,
        mlir.lower_fun(fn, multiple_results=False),
        platform=platform,
        inline=False)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to a numeric dtype first: x.astype(jnp.float32)
2. For bool masks use jnp.cumsum which promotes, or lax.cumsum on x.astype(np.int32)
3. Fix upstream logic so a numeric array reaches the primitive

Example fix

// before
jax.lax.cumsum(bool_mask)
// after
jax.lax.cumsum(bool_mask.astype(jnp.int32))

if not jnp.issubdtype(x.dtype, jnp.number):
    x = x.astype(jnp.float32)

def is_numeric_array(x) -> bool:
    import numpy as np, jax.numpy as jnp
    return jnp.issubdtype(x.dtype, np.number)

When it happens

Trigger: Calling jax.lax.cumlogsumexp (or cumsum in stricter versions) on a boolean or non-numeric array, e.g. cumsum(jnp.array([True, False])).

Common situations: Mask arrays (bool) fed to a cumulative op by accident; pipelines where a preceding comparison produces bool that then reaches cumsum; object/str arrays leaking from data preprocessing.

Related errors

• XLA operations do not allow negative axes
• axis {} is out of bounds for array of shape {}
• logical reduction requires operand dtype bool or int, got {o
• logaddexp2 requires floating-point or complex inputs; got {x
• index arguments to dynamic_update_slice must be integers of

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