jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

Raised in the SparseCore reduce lowering: hardware reductions operate along the innermost axis only, so the axes parameter must be exactly (0,); any other axis specification raises. The input at fault is the axes argument of a reduce inside a SparseCore kernel.

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.
• Cannot bitcast a ()-shaped array to a dtype with a…
• Cannot bitcast from ( bits) to ( bits), because * is not…
• cummax: x= must be rank 1
• duplicate value in 'axes' of reduction

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

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