jax-ml/jax · error · NotImplementedError

reductions require axes to be (0,) on SparseCore, but got {a

Error message

reductions require axes to be (0,) on SparseCore, but got {axes}.

What it means

SparseCore reduction lowering only supports reducing over axis 0 (the subcore's vector dimension). Reductions over any other axes tuple raise NotImplementedError.

Source

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

    )
  if sign_bit_vec is not None:  # Flip the sign bit back
    return arith.xori(result, sign_bit_vec)
  return result


sc_lowering.register_lowering_rule(masked_cummax_p)(
    functools.partial(_masked_cumop_lowering_rule, reduction_kind="max"))
sc_lowering.register_lowering_rule(masked_cummin_p)(
    functools.partial(_masked_cumop_lowering_rule, reduction_kind="min"))
sc_lowering.register_lowering_rule(masked_cumsum_p)(
    functools.partial(_masked_cumop_lowering_rule, reduction_kind="sum"))


def _reduce_op_lowering_rule(ctx: sc_lowering.LoweringRuleContext, x, axes,
                             *, reduction_kind, out_sharding=None):
  del out_sharding  # Unused.
  if axes != (0,):
    raise NotImplementedError(
        f"reductions require axes to be (0,) on SparseCore, but got {axes}.")
  vec_dim = ctx.avals_in[0].shape[0]
  i1t = ir.IntegerType.get_signless(1)
  c1 = arith.constant(i1t, ir.IntegerAttr.get(i1t, 1))
  x_shp = ctx.avals_in[0].shape
  c1v = vector.broadcast(ir.VectorType.get(x_shp, c1.type), c1)
  return vector.extract(
      _masked_cumop_lowering_rule(ctx, x, c1v, reduction_kind=reduction_kind),
      [], [vec_dim - 1])

sc_lowering.register_lowering_rule(
    lax.reduce_max_p, kernel_types=[tpu_core.CoreType.SC_VECTOR_SUBCORE])(
    functools.partial(_reduce_op_lowering_rule, reduction_kind="max"))
sc_lowering.register_lowering_rule(
    lax.reduce_min_p, kernel_types=[tpu_core.CoreType.SC_VECTOR_SUBCORE])(
    functools.partial(_reduce_op_lowering_rule, reduction_kind="min"))
sc_lowering.register_lowering_rule(
    lax.reduce_sum_p, kernel_types=[tpu_core.CoreType.SC_VECTOR_SUBCORE])(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose the array so the reduced dimension is axis 0, reduce, then transpose back if needed
2. Restructure to only reduce across the vector dimension
3. Compute multi-axis reductions as a sequence of axis-0 reductions

Example fix

// before
y = sc_reduce(x, axes=(1,))

// after
y = sc_reduce(x.swapaxes(0, 1), axes=(0,))  # then swap result back as needed

assert axes == (0,), f'SparseCore reductions only support axes=(0,), got {axes}'

def sc_reduce_axes_ok(axes) -> bool:
    return tuple(axes) == (0,)

When it happens

Trigger: Binding a reduction primitive (sum/max/min over a BlockLayout array) with axes=(1,) or axes=(0,1) instead of (0,).

Common situations: Porting lax-style multi-axis reductions into an SC kernel; transposing data so the reduced dim is not axis 0.

Related errors

• `pallas` reduce operations only support one reduce axis.
• duplicate value in 'axes' of reduction: {axes}
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Only single axis reduction supported

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