jax-ml/jax · error · NotImplementedError

Only single axis reduction supported

Error message

Only single axis reduction supported

What it means

The Mosaic reduce_index helper only lowers reductions over a single axis. If argmin/argmax is requested over multiple axes simultaneously (len(axes) != 1), it raises NotImplementedError('Only single axis reduction supported').

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3568

  elif jnp.issubdtype(aval_out.dtype, jnp.floating):
    return arith.minimumf(x, y)
  raise NotImplementedError(aval_out.dtype)

def _reduce_index_helper(
    ctx: LoweringRuleContext, x, axes, index_dtype, reduction_kind):
  (x_aval,) = ctx.avals_in
  (out_aval,) = ctx.avals_out
  if (x_aval.dtype, index_dtype) not in (
      (jnp.float32, jnp.int32),
      (jnp.bfloat16, jnp.int16),
      (jnp.bfloat16, jnp.int32),
  ):
    raise NotImplementedError(
        f"Unsupported combination of input dtype ({x_aval.dtype}) and"
        f" index_dtype ({index_dtype}) for reduce_index"
    )
  if len(axes) != 1:
    raise NotImplementedError("Only single axis reduction supported")

  axis = axes[0]
  # TODO(b/460843515): Support 1D inputs in Mosaic.
  is_1d = len(x_aval.shape) == 1
  if is_1d:
    x = vector.shape_cast(
        ctx.aval_to_ir_type(
            jax_core.ShapedArray((1, *x_aval.shape), x_aval.dtype)
        ),
        x,
    )
    axis += 1
    out_shape = (1, *out_aval.shape)
  else:
    out_shape = out_aval.shape

  native_dtype = jnp.int16 if x_aval.dtype == jnp.bfloat16 else jnp.int32
  native_out_type = ctx.aval_to_ir_type(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten explicitly then reduce once: jnp.argmax(x.reshape(-1)) → but prefer per-axis: reshape to put the target axis last and use axis=-1
2. Loop over axes and combine results manually
3. Compute the global argmax outside the kernel
4. Restructure so the kernel reduces a single blocking axis

Example fix

// before
i = jnp.argmax(x, axis=None)  # in kernel
// after
flat = x.reshape(-1)
i = jnp.argmax(flat)  # single logical axis; or compute outside kernel

def single_axis_only(axes):
    return axes is None or len(axes) == 1 or isinstance(axes, int)

When it happens

Trigger: Calling jnp.argmax(x, axis=None) or a multi-axis argmax inside a Pallas kernel, which collapses all axes into one multi-axis reduction.

Common situations: Using axis=None argmax on a 2D block inside a kernel; refactoring NumPy code that used global argmax.

Related errors

• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsupported combination of input dtype ({x_aval.dtype}) and
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

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