{"record":{"id":"3fdcdc84feb95d28","repo":"jax-ml/jax","slug":"only-gathers-along-the-two-minormost-dimensions-su","errorCode":null,"errorMessage":"Only gathers along the two minormost dimensions supported on TC","messagePattern":"Only gathers along the two minormost dimensions supported on TC","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3408,"sourceCode":"  if (\n      slice_sizes == (1,) * rank\n      and mode\n      in (\n          lax.GatherScatterMode.FILL_OR_DROP,\n          lax.GatherScatterMode.PROMISE_IN_BOUNDS,\n      )\n      and not offset_dims\n      and collapsed_slice_dims == start_index_map\n      and operand_batching_dims == start_indices_batching_dims\n      and len(collapsed_slice_dims) == 1\n      and len(operand_batching_dims) == rank - 1\n  ):\n    (axis,) = collapsed_slice_dims\n    if (\n        ctx.lowering_context.kernel_type == tpu_core.CoreType.TC\n        and axis < rank - 2\n    ):\n      raise NotImplementedError(\n          \"Only gathers along the two minormost dimensions supported on TC\"\n      )\n    return tpu.dynamic_gather(x, recovered_indices, [axis])\n  raise NotImplementedError(\"Unsupported gather\")\n\n\n@register_lowering_rule(lax.transpose_p)\ndef _transpose_lowering_rule(ctx: LoweringRuleContext, x, *, permutation):\n  out_type = ctx.aval_to_ir_type(ctx.avals_out[0])\n  return tpu.transpose(out_type, x, permutation)\n\n\ndef _bcast(\n    x: ir.Value | object,\n    y: ir.Value | object,\n    x_aval: ShapedAbstractValue,\n    y_aval: ShapedAbstractValue,\n    out_aval: ShapedAbstractValue,","sourceCodeStart":3390,"sourceCodeEnd":3426,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3390-L3426","documentation":"On TPU TensorCore cores, Mosaic can only lower dynamic gathers that operate on one of the two minormost (last two) dimensions of the tensor. A gather whose collapsed axis is further from the minor end falls back to an unsupported path on TC hardware and raises this NotImplementedError during kernel compilation.","triggerScenarios":"Running a Pallas kernel on a TensorCore-targeted compilation where take_along_axis/dynamic_gather is along an axis with index < rank-2 (e.g. axis 0 of a 4D tensor).","commonSituations":"Writing a kernel that gathers along batch or heads dimension of a 4D activation tensor; code that worked on TPU v4/v5e scalar/vector paths failing when scheduled onto TC subcores; attention-style index selects on non-minor axes.","solutions":["Transpose the tensor so the gather axis becomes one of the last two dimensions, gather, then transpose back","Rework the kernel layout (e.g. use BlockMapping to put the gathered axis minormost)","Do the gather outside the kernel as a pre-processing step","Target a non-TC core type if the kernel doesn't need TensorCore ops"],"exampleFix":"// before\nout = jnp.take_along_axis(x, idx, axis=0)  # x is 4D, axis 0 not minormost\n// after\nx_t = jnp.transpose(x, (1,2,3,0))\nout_t = jnp.take_along_axis(x_t, idx_t, axis=-1)\nout = jnp.transpose(out_t, (3,0,1,2))","handlingStrategy":"validation","validationCode":"def gather_ok_on_tc(rank, axis):\n    return axis >= rank - 2","typeGuard":null,"tryCatchPattern":"catch NotImplementedError at compile time and restructure to transpose the gather axis minormost","preventionTips":["Keep gather axes minormost in kernel memory layouts","Document per-core-type assumptions in kernel docstrings","Test on the target core type in CI"],"tags":["jax","pallas","tpu","tensorcore","gather"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}