{"record":{"id":"79c3613a50c73d53","repo":"jax-ml/jax","slug":"cummax-x-x-aval-must-be-rank-1","errorCode":null,"errorMessage":"cummax: x={x.aval} must be rank 1","messagePattern":"cummax: x=(.+?) must be rank 1","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":664,"sourceCode":"    lax.reduce_sum_p, kernel_types=[tpu_core.CoreType.SC_VECTOR_SUBCORE])(\n    functools.partial(_reduce_op_lowering_rule, reduction_kind=\"sum\"))\n\n\ndef cummax(x: jax.Array, *, mask: jax.Array | None = None) -> jax.Array:\n  \"\"\"Returns the cumulative max of the array along its innermost axis.\n\n  Elements from `x` will pass through directly to the result until the first\n  valid value is encountered (`mask[i] == True`). If you would like to specify\n  a default value for such elements instead, write\n  `x = jnp.where(mask, x, default_value)` before or after calling this function.\n\n  Args:\n    x: An array of integers or floats.\n    mask: An optional array of booleans, which specifies which elements of `x`\n      are eligible for the max. If `None`, all elements are eligible.\n  \"\"\"\n  if x.ndim != 1:\n    raise NotImplementedError(f\"cummax: x={x.aval} must be rank 1\")\n  if mask is None:\n    mask = lax.full(x.shape, True)\n  return masked_cummax_p.bind(x, mask)\n\n\ndef cummin(x: jax.Array, *, mask: jax.Array | None = None) -> jax.Array:\n  \"\"\"Returns the cumulative min of the array along its innermost axis.\n\n  Elements from `x` will pass through directly to the result until the first\n  valid value is encountered (`mask[i] == True`). If you would like to specify\n  a default value for such elements instead, write\n  `x = jnp.where(mask, x, default_value)` before or after calling this function.\n\n  Args:\n    x: An array of integers or floats.\n    mask: An optional array of booleans, which specifies which elements of `x`\n      are eligible for the min. If `None`, all elements are eligible.\n  \"\"\"","sourceCodeStart":646,"sourceCodeEnd":682,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L646-L682","documentation":"The public SC cummax wrapper only accepts rank-1 arrays; higher-rank inputs raise NotImplementedError because the hardware cummax operates on a single vector.","triggerScenarios":"cummax(x) with x.ndim > 1 (e.g. shape (B, N)).","commonSituations":"Applying per-batch cumulative max without vmap; feeding attention matrices (rank 2+) directly.","solutions":["Use jax.vmap(cummax, in_axes=1) (or 0) to map over the extra dimension","Reshape/loop over the leading dims and call cummax per row","Check x.ndim before calling in generic code"],"exampleFix":"// before\ny = cummax(x)  # x.shape = (B, N)\n\n// after\ny = jax.vmap(cummax, in_axes=1, out_axes=1)(x)","handlingStrategy":"validation","validationCode":"assert x.ndim == 1, f'cummax requires rank-1 input, got {x.ndim}'\n# or: y = jax.vmap(cummax, in_axes=1, out_axes=1)(x) for rank-2","typeGuard":"def is_rank1(x) -> bool:\n    return getattr(x, 'ndim', None) == 1","tryCatchPattern":null,"preventionTips":["vmap over batch dims instead of passing rank-2 inputs","Add ndim checks in generic code paths","Remember SC vector ops are rank-1 by design"],"tags":["jax","pallas","sparsecore","cummax","rank"],"backgroundTag":"operand-rank-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}