{"record":{"id":"f63847ca188fdbf2","repo":"jax-ml/jax","slug":"unsupported-gather","errorCode":null,"errorMessage":"Unsupported gather","messagePattern":"Unsupported gather","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3412,"sourceCode":"          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,\n    dynamic_shape_replacement_fn: DynamicShapeReplacementFn,\n) -> tuple[ir.Value, ir.Value]:\n  x_dtype = x_aval.dtype\n  y_dtype = y_aval.dtype","sourceCodeStart":3394,"sourceCodeEnd":3430,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3394-L3430","documentation":"Catch-all failure in the Mosaic lax.gather lowering rule: the gather's dimension numbers did not match any of the recognized patterns (offset-only collapsed dims single-axis, or start_index_map gathers). The TPU backend only implements a narrow subset of JAX's general gather semantics.","triggerScenarios":"Calling lax.gather (directly or via jnp.ix_/fancy indexing traced into a kernel) with collapsed_slice_dims or start_index_map combinations outside the two supported cases in the lowering rule.","commonSituations":"Using arbitrary fancy indexing inside jax.pallas kernels; mixing gather with batched index arrays the lowering doesn't recognize; relying on NumPy-style advanced indexing semantics in kernel code.","solutions":["Re-express the operation as jnp.take_along_axis on a single axis","Compute the gather outside the Pallas kernel","Rewrite using explicit tpu primitives or manual index arithmetic with tpu.load","Upgrade JAX — gather coverage in Mosaic improves over versions"],"exampleFix":"// before\nvals = lax.gather(x, index_cm, dnums, slice_sizes)\n// after\nvals = x[flat_idx]  # computed outside kernel, passed as input","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"except NotImplementedError: fall back to computing the gather outside the Pallas kernel with plain jnp indexing","preventionTips":["Prefer take_along_axis over lax.gather in kernels","Keep fancy indexing out of Pallas kernels"],"tags":["jax","pallas","tpu","gather","not-implemented"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}