{"record":{"id":"228a70fe62fcc770","repo":"jax-ml/jax","slug":"does-not-accept-dtype-accepted-dtypes-are-s","errorCode":null,"errorMessage":"{} does not accept dtype {}. Accepted dtypes are subtypes of number.","messagePattern":"(.+?) does not accept dtype (.+?)\\. Accepted dtypes are subtypes of number\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":3094,"sourceCode":"    return x\n  padding = [(0, 0)] * x.ndim\n  padding[axis] = (0, n - 1) if reverse else (n - 1, 0)\n  strides = [1] * x.ndim\n  window_dims = [1] * x.ndim\n  window_dims[axis] = n\n  return window_reduce(x, window_dims, strides, padding)\n\n\ndef _cumred_batch_rule(prim, batched_args, batch_dims, *, axis: int,\n                       reverse: bool):\n  operand, = batched_args\n  bdim, = batch_dims\n  axis = axis if axis < bdim else axis + 1\n  return prim.bind(operand, axis=axis, reverse=reverse), bdim\n\ndef _cumred_dtype_rule(name, operand, *args, **kw):\n  if not dtypes.issubdtype(operand.dtype, np.number):\n    raise TypeError(\"{} does not accept dtype {}. Accepted dtypes are subtypes \"\n                    \"of number.\".format(name, np.dtype(operand.dtype).name))\n  return operand.dtype\n\n\ndef _cumulative_reduction_primitive(name, reduce_fn, reduce_window_fn):\n  reducer_p = lax.standard_primitive(\n    _cumred_shape_rule, partial(_cumred_dtype_rule, name),\n    name, sharding_rule=_cumred_sharding_rule,\n    vma_rule=partial(core.standard_vma_rule, name))\n  batching.primitive_batchers[reducer_p] = partial(_cumred_batch_rule,\n                                                   reducer_p)\n\n  def register_lowering(fn, platform=None):\n    mlir.register_lowering(\n        reducer_p,\n        mlir.lower_fun(fn, multiple_results=False),\n        platform=platform,\n        inline=False)","sourceCodeStart":3076,"sourceCodeEnd":3112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L3076-L3112","documentation":"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.","triggerScenarios":"Calling jax.lax.cumlogsumexp (or cumsum in stricter versions) on a boolean or non-numeric array, e.g. cumsum(jnp.array([True, False])).","commonSituations":"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.","solutions":["Cast to a numeric dtype first: x.astype(jnp.float32)","For bool masks use jnp.cumsum which promotes, or lax.cumsum on x.astype(np.int32)","Fix upstream logic so a numeric array reaches the primitive"],"exampleFix":"// before\njax.lax.cumsum(bool_mask)\n// after\njax.lax.cumsum(bool_mask.astype(jnp.int32))","handlingStrategy":"type-guard","validationCode":"if not jnp.issubdtype(x.dtype, jnp.number):\n    x = x.astype(jnp.float32)","typeGuard":"def is_numeric_array(x) -> bool:\n    import numpy as np, jax.numpy as jnp\n    return jnp.issubdtype(x.dtype, np.number)","tryCatchPattern":null,"preventionTips":["Cast bool masks to int/float before cumulative ops","Watch for comparison ops producing bool arrays upstream"],"tags":["jax","lax","dtype","cumsum"],"backgroundTag":"invalid-dtype-for-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}