jax-ml/jax · error · NotImplementedError

cummax: x= must be rank 1

Error message

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

What it means

Validation in cummax on Pallas SparseCore: cumulative max runs along the innermost axis of a rank-1 vector, so an input whose aval is not rank 1 raises. The input at fault is the x argument of cummax; flatten or slice to a 1D vector first.

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

• Cannot bitcast a ()-shaped array to a dtype with a…
• Cannot bitcast from ( bits) to ( bits), because * is not…
• Expected -dim 'key' tensor for MQA. Instead got a -dim one.
• Expected metadata to be 3-dimensional (M, K // 4, 2), but…
• Gather only supports loading from VMEM, got

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

Appendix: 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)