{"record":{"id":"504c7b3fcf54df3e","repo":"jax-ml/jax","slug":"pallas-reduce-operations-only-support-one-reduce","errorCode":null,"errorMessage":"`pallas` reduce operations only support one reduce axis.","messagePattern":"`pallas` reduce operations only support one reduce axis\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2542,"sourceCode":"\ntriton_lowering_rules[lax.reduce_max_p] = functools.partial(\n    _reduce_lowering, jnp.maximum\n)\ntriton_lowering_rules[lax.reduce_min_p] = functools.partial(\n    _reduce_lowering, jnp.minimum\n)\ntriton_lowering_rules[lax.reduce_sum_p] = functools.partial(\n    _reduce_lowering, jnp.add\n)\n\n\ndef _argreduce_lowering(\n    body, ctx: LoweringRuleContext, a, *, axes, index_dtype\n):\n  if index_dtype != jnp.int32:\n    raise ValueError(\"`index_type` must be i32.\")\n  if len(axes) != 1:\n    raise ValueError(\"`pallas` reduce operations only support one reduce axis.\")\n  [axis] = axes\n  [a_aval] = ctx.avals_in\n  index = _make_range(0, a_aval.shape[axis])\n  if len(a_aval.shape) > 1:\n    # Broadcast index across the non-reduced axes\n    for i in range(len(a_aval.shape)):\n      if i != axis:\n        index = _expand_dims(index, i)\n    index = _bcast_to(index, a_aval.shape)\n  ctx = ctx.replace(avals_in=[a_aval, a_aval.update(dtype=jnp.dtype(jnp.int32))])\n  _, indices = _reduction_lowering(body, ctx, (a, index), axes=axes)\n  return indices\n\n\ndef _reduce_argmax_combine(left, right):\n  value1, index1 = left\n  value2, index2 = right\n  gt = value1 > value2","sourceCodeStart":2524,"sourceCodeEnd":2560,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2524-L2560","documentation":"Pallas triton reduce lowerings only implement a single reduction axis. If lax/arg-reduce is called with multiple axes (e.g. reduce over two dimensions at once), the lowering raises ValueError immediately.","triggerScenarios":"Using jnp.sum(x, axis=(0, 1)) or argmax with multiple axes inside a pallas kernel body; reducing a whole block with axis=None expanded to multiple axes.","commonSituations":"Porting vectorized numpy-style reductions to block-level kernels; assuming XLA multi-axis reduce semantics carry into Mosaic.","solutions":["Split into sequential single-axis reductions: jnp.sum(jnp.sum(x, axis=0), axis=0)","Reshape to merge axes before reducing: x.reshape(-1, ...) then reduce over one axis","Perform multi-axis reductions outside the pallas kernel"],"exampleFix":"# before\ntotal = jnp.sum(block, axis=(0, 1))\n\n# after\ntotal = jnp.sum(block.reshape(-1, block.shape[-1]), axis=0)","handlingStrategy":"validation","validationCode":"assert isinstance(axis, int) or (isinstance(axis, tuple) and len(axis) == 1), \\\n    'pallas reductions support exactly one axis'","typeGuard":"def single_axis(axis) -> bool:\n    return isinstance(axis, int) or (isinstance(axis, (tuple, list)) and len(axis) == 1)","tryCatchPattern":null,"preventionTips":["Flatten axes with reshape before reducing inside kernels"],"tags":["jax","pallas","triton","reduction","axes"],"backgroundTag":"unsupported-multi-axis-reduction","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}