jax-ml/jax · error · NotImplementedError

Unsupported combination of input dtype ({x_aval.dtype}) and

Error message

Unsupported combination of input dtype ({x_aval.dtype}) and index_dtype ({index_dtype}) for reduce_index

What it means

The reduce_index lowering (argmin/argmax-style reductions that return indices) only accepts the (input dtype, index dtype) pairs (float32,int32), (bfloat16,int16), (bfloat16,int32). Any other combination raises NotImplementedError at lowering time.

Source

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

  (aval_out,) = ctx.avals_out
  if jnp.issubdtype(aval_out.dtype, jnp.signedinteger):
    return arith.minsi(x, y)
  elif jnp.issubdtype(aval_out.dtype, jnp.unsignedinteger):
    return arith.minui(x, y)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the operand to jnp.float32 (or bfloat16) before argmax/argmin
2. Pass index_dtype explicitly: lax.argmax(x, axis, index_dtype=jnp.int32)
3. Compute indices outside the kernel
4. Upgrade JAX — reduce_index dtype coverage has been expanding

Example fix

// before
i = lax.argmax(x, axis=0, index_dtype=jnp.uint16)
// after
i = lax.argmax(x.astype(jnp.float32), axis=0, index_dtype=jnp.int32)

import jax.numpy as jnp
SUPPORTED = {(jnp.float32, jnp.int32), (jnp.bfloat16, jnp.int16), (jnp.bfloat16, jnp.int32)}
def reduce_index_ok(x_dtype, idx_dtype):
    return any(jnp.dtype(x_dtype) == a and jnp.dtype(idx_dtype) == b for a, b in SUPPORTED)

When it happens

Trigger: Calling jnp.argmin/jnp.argmax (or lax.argmin/argmax) inside a Pallas kernel where the operand is not float32/bfloat16 or the index_dtype parameter is not int16/int32 (e.g. default uint16/uint32 index dtype on some paths, or float16 input).

Common situations: Using argmax over float16 logits inside a kernel (common in inference kernels); lax.argmax with index_dtype=jnp.uint16; mixing bfloat16 with an int64 index request.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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