jax-ml/jax · error · ValueError

`pallas` reduce operations only support one reduce axis.

Error message

`pallas` reduce operations only support one reduce axis.

What it means

Pallas triton reduce lowerings only implement a single reduction axis. If lax/arg-reduce is called with multiple axes (e.g. reduce over two dimensions at once), the lowering raises ValueError immediately.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2542

triton_lowering_rules[lax.reduce_max_p] = functools.partial(
    _reduce_lowering, jnp.maximum
)
triton_lowering_rules[lax.reduce_min_p] = functools.partial(
    _reduce_lowering, jnp.minimum
)
triton_lowering_rules[lax.reduce_sum_p] = functools.partial(
    _reduce_lowering, jnp.add
)


def _argreduce_lowering(
    body, ctx: LoweringRuleContext, a, *, axes, index_dtype
):
  if index_dtype != jnp.int32:
    raise ValueError("`index_type` must be i32.")
  if len(axes) != 1:
    raise ValueError("`pallas` reduce operations only support one reduce axis.")
  [axis] = axes
  [a_aval] = ctx.avals_in
  index = _make_range(0, a_aval.shape[axis])
  if len(a_aval.shape) > 1:
    # Broadcast index across the non-reduced axes
    for i in range(len(a_aval.shape)):
      if i != axis:
        index = _expand_dims(index, i)
    index = _bcast_to(index, a_aval.shape)
  ctx = ctx.replace(avals_in=[a_aval, a_aval.update(dtype=jnp.dtype(jnp.int32))])
  _, indices = _reduction_lowering(body, ctx, (a, index), axes=axes)
  return indices


def _reduce_argmax_combine(left, right):
  value1, index1 = left
  value2, index2 = right
  gt = value1 > value2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Split into sequential single-axis reductions: jnp.sum(jnp.sum(x, axis=0), axis=0)
Reshape to merge axes before reducing: x.reshape(-1, ...) then reduce over one axis
Perform multi-axis reductions outside the pallas kernel

Example fix

# before
total = jnp.sum(block, axis=(0, 1))

# after
total = jnp.sum(block.reshape(-1, block.shape[-1]), axis=0)

Defensive patterns

Strategy: validation

Validate before calling

assert isinstance(axis, int) or (isinstance(axis, tuple) and len(axis) == 1), \
    'pallas reductions support exactly one axis'

Type guard

def single_axis(axis) -> bool:
    return isinstance(axis, int) or (isinstance(axis, (tuple, list)) and len(axis) == 1)

Prevention

Flatten axes with reshape before reducing inside kernels

When it happens

Trigger: Using jnp.sum(x, axis=(0, 1)) or argmax with multiple axes inside a pallas kernel body; reducing a whole block with axis=None expanded to multiple axes.

Common situations: Porting vectorized numpy-style reductions to block-level kernels; assuming XLA multi-axis reduce semantics carry into Mosaic.

Related errors

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