{"record":{"id":"48172ddcb4d1e253","repo":"jax-ml/jax","slug":"only-single-axis-reduction-supported","errorCode":null,"errorMessage":"Only single axis reduction supported","messagePattern":"Only single axis reduction supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3568,"sourceCode":"  elif jnp.issubdtype(aval_out.dtype, jnp.floating):\n    return arith.minimumf(x, y)\n  raise NotImplementedError(aval_out.dtype)\n\ndef _reduce_index_helper(\n    ctx: LoweringRuleContext, x, axes, index_dtype, reduction_kind):\n  (x_aval,) = ctx.avals_in\n  (out_aval,) = ctx.avals_out\n  if (x_aval.dtype, index_dtype) not in (\n      (jnp.float32, jnp.int32),\n      (jnp.bfloat16, jnp.int16),\n      (jnp.bfloat16, jnp.int32),\n  ):\n    raise NotImplementedError(\n        f\"Unsupported combination of input dtype ({x_aval.dtype}) and\"\n        f\" index_dtype ({index_dtype}) for reduce_index\"\n    )\n  if len(axes) != 1:\n    raise NotImplementedError(\"Only single axis reduction supported\")\n\n  axis = axes[0]\n  # TODO(b/460843515): Support 1D inputs in Mosaic.\n  is_1d = len(x_aval.shape) == 1\n  if is_1d:\n    x = vector.shape_cast(\n        ctx.aval_to_ir_type(\n            jax_core.ShapedArray((1, *x_aval.shape), x_aval.dtype)\n        ),\n        x,\n    )\n    axis += 1\n    out_shape = (1, *out_aval.shape)\n  else:\n    out_shape = out_aval.shape\n\n  native_dtype = jnp.int16 if x_aval.dtype == jnp.bfloat16 else jnp.int32\n  native_out_type = ctx.aval_to_ir_type(","sourceCodeStart":3550,"sourceCodeEnd":3586,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3550-L3586","documentation":"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').","triggerScenarios":"Calling jnp.argmax(x, axis=None) or a multi-axis argmax inside a Pallas kernel, which collapses all axes into one multi-axis reduction.","commonSituations":"Using axis=None argmax on a 2D block inside a kernel; refactoring NumPy code that used global argmax.","solutions":["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","Loop over axes and combine results manually","Compute the global argmax outside the kernel","Restructure so the kernel reduces a single blocking axis"],"exampleFix":"// before\ni = jnp.argmax(x, axis=None)  # in kernel\n// after\nflat = x.reshape(-1)\ni = jnp.argmax(flat)  # single logical axis; or compute outside kernel","handlingStrategy":"validation","validationCode":"def single_axis_only(axes):\n    return axes is None or len(axes) == 1 or isinstance(axes, int)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid axis=None argmax inside kernels","Flatten or loop over axes explicitly"],"tags":["jax","pallas","tpu","argmax","reduction"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}