jax-ml/jax · error · NotImplementedError

cummax: x={x.aval} must be rank 1

Error message

cummax: x={x.aval} must be rank 1

What it means

The public SC cummax wrapper only accepts rank-1 arrays; higher-rank inputs raise NotImplementedError because the hardware cummax operates on a single vector.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:664

    lax.reduce_sum_p, kernel_types=[tpu_core.CoreType.SC_VECTOR_SUBCORE])(
    functools.partial(_reduce_op_lowering_rule, reduction_kind="sum"))


def cummax(x: jax.Array, *, mask: jax.Array | None = None) -> jax.Array:
  """Returns the cumulative max of the array along its innermost axis.

  Elements from `x` will pass through directly to the result until the first
  valid value is encountered (`mask[i] == True`). If you would like to specify
  a default value for such elements instead, write
  `x = jnp.where(mask, x, default_value)` before or after calling this function.

  Args:
    x: An array of integers or floats.
    mask: An optional array of booleans, which specifies which elements of `x`
      are eligible for the max. If `None`, all elements are eligible.
  """
  if x.ndim != 1:
    raise NotImplementedError(f"cummax: x={x.aval} must be rank 1")
  if mask is None:
    mask = lax.full(x.shape, True)
  return masked_cummax_p.bind(x, mask)


def cummin(x: jax.Array, *, mask: jax.Array | None = None) -> jax.Array:
  """Returns the cumulative min of the array along its innermost axis.

  Elements from `x` will pass through directly to the result until the first
  valid value is encountered (`mask[i] == True`). If you would like to specify
  a default value for such elements instead, write
  `x = jnp.where(mask, x, default_value)` before or after calling this function.

  Args:
    x: An array of integers or floats.
    mask: An optional array of booleans, which specifies which elements of `x`
      are eligible for the min. If `None`, all elements are eligible.
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.vmap(cummax, in_axes=1) (or 0) to map over the extra dimension
2. Reshape/loop over the leading dims and call cummax per row
3. Check x.ndim before calling in generic code

Example fix

// before
y = cummax(x)  # x.shape = (B, N)

// after
y = jax.vmap(cummax, in_axes=1, out_axes=1)(x)

assert x.ndim == 1, f'cummax requires rank-1 input, got {x.ndim}'
# or: y = jax.vmap(cummax, in_axes=1, out_axes=1)(x) for rank-2

def is_rank1(x) -> bool:
    return getattr(x, 'ndim', None) == 1

When it happens

Trigger: cummax(x) with x.ndim > 1 (e.g. shape (B, N)).

Common situations: Applying per-batch cumulative max without vmap; feeding attention matrices (rank 2+) directly.

Related errors

• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor
• {self} should have a different axis name from the TensorCore

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