[zsun6]

aten.gather.default, aten.scatter.src and aten.scatter.value had no MLX handler. They're core ATen so to_edge doesn't decompose them, and any graph with one of them splits at that node: the op runs on the portable CPU kernel with a delegate handoff on each side.

Subgraphs from this repo, before/after (to_edge_transform_and_lower with MLXPartitioner):

subgraph	before	after
Sortformer export-safe rel_shift (examples/models/sortformer/export_sortformer.py:164, one per attention layer), x (1,8,64,127)	1 segment, gather on CPU	1 segment, nothing on CPU
Muse Glimmer sampling_probabilities (examples/models/muse-glimmer/model/dflash_token_sampler.py:63), logits (4,512)	3 segments; scatter.src, sort.stable on CPU	2 segments; sort.stable on CPU
Muse Glimmer verify_speculative (same file, R=3, V=512)	5 segments; 2x gather, scatter.src, scatter.value, sort.stable, 2x sum on CPU	3 segments; sort.stable, 2x sum on CPU
top-2 MoE router as in examples/models/llama/llama_transformer.py:156 (scores.gather(1, idx)) and HF gpt_oss / qwen3_vl_moe (zeros_like(logits).scatter_(1, idx, w)), (16,256) -> 8 experts	2 segments; gather, scatter.src on CPU	1 segment, nothing on CPU

(sort.stable isn't core ATen and needs _core_aten_ops_exception_list; sum over bool is a separate gap. Both out of scope here.)

How

gather lowers onto the existing TakeAlongAxisNode. aten allows index.size(d) <= self.size(d) on the non-gather axes and only reads that leading block, while mlx::take_along_axis broadcasts, so the handler first narrows self to the index shape on those axes with SliceNodes (nothing is emitted when the sizes already agree, including a shared symbolic dim; a static index size is sliced to even if the input dim is symbolic). The output has the index shape, so the index can be longer than self along dim for free.

scatter adds a PutAlongAxisNode (appended at the end of the OpNode union, so existing .pte files still load) and an exec_put_along_axis that calls mlx::put_along_axis. Two aten freedoms need handling: src may be larger than index on any axis (only its leading index.shape block is read), which the handler narrows like gather; and index may be smaller than self on the non-scatter axes, which the kernel handles by scattering into the leading index-shaped block of self and writing it back with slice_update — broadcasting there would write every row instead of the first index.size(d). A scalar value becomes a 0-D constant that MLX broadcasts, so scatter.value shares the node. int64/float64 inputs are declined at partition time because ScatterAxis::eval_gpu has no kernel for 8-byte element types. Duplicate indices along dim are unspecified in aten; MLX keeps one of the writes. Indices aren't bounds-checked in the delegate, same as the existing take_along_axis/scatter_add lowerings.

Tests

Adds gather and scatter to backends/mlx/test/test_ops.py, 32 configs: 1D-4D, every axis and negative dim, index smaller than the input on the other axes and longer along dim, src larger than index, scalar value (including a float written into an int32 self), fp16/bf16/int64 inputs with int64 indices, a gather from a transposed view, and a dynamic batch dim (shared by the operands, and on the input alone against a static index) exported at one size and run at another. Both ops are pure data movement, so the comparison is exact (rtol = atol = 0).

python -m executorch.backends.mlx.test.run_all_tests --rebuild gather scatter   # 32 passed, 0 failed
python -m executorch.backends.mlx.test.run_all_tests -j4 --clean-after           # 982 passed, 0 failed

M3 Pro, macOS 14.1, Xcode 15.3, torch 2.14.0, MLX 1f8e74e3. Only schema.fbs is committed; the generated files come from the build, as in #23051.

Written with AI assistance (Claude Code). I reviewed the design and the diff, ran the full MLX op suite locally and take responsibility for the change.

cc @nil-is-all @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23398

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures

As of commit ca3cc1c with merge base a6171a6 ():

NEW FAILURES - The following jobs have failed:

• Test QNN Backend / test-qnn / test-backend-linux (qnn, models) / linux-job (gh)
  curl: (22) The requested URL returned error:
• Test QNN Backend / test-qnn / test-backend-linux (qnn, operators) / linux-job (gh)
  curl: (22) The requested URL returned error:
• Test QNN Windows MSVC build / build-qnn-windows-arm64 (gh)
• Test QNN Windows MSVC build / build-qnn-windows-x64 / windows-job (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @zsun6!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: zsun6 / name: Zhongrui Sun (e6b6332)

[nil-is-all]

Hi @zsun6, thanks for the PR! Could you rebase and ping for review when ready? This would ensure the CI workflows focus on this PR's code changes post some fixes from our end.

[zsun6]

@nil-is-all Rebased onto main in 0c37c45. The gather/scatter export tests and C++ syntax checks pass. @metascroy Could you take another look and approve MLX CI?

[zsun6]

@pytorchbot label "release notes: ops & kernels"

[metascroy]

@claude review this code. Do the test cases cover corner cases or odd configs? Doe the node fbs shema / Node closely follow the MLX op's contract like other op handlers (not aten). Does the schema hand dynamism appropriately, using Vid where it could

[claude]

Claude finished @metascroy's task in 2m 42s —— View job

---

PutAlongAxisNode's schema matches mlx::put_along_axis, but exec_put_along_axis builds aten semantics into the runtime, _slice_to_shape turns away symbolic sizes it could handle with a Vid, and the tests leave a few dynamic-shape and decline cases unchecked. Needs Discussion.

Detailed review

PR Review: #23398

Summary

The PR lowers aten.gather onto TakeAlongAxisNode and aten.scatter.{src,value} onto a new PutAlongAxisNode. The node's fields match MLX's put_along_axis(a, indices, values, axis). The C++ executor does not: it carries aten's narrowing logic itself. On dynamic shapes, the _slice_to_shape docstring says a slice stop is fixed at export time. That is wrong (SliceNode.stop is an IntOrVid), so the handler declines graphs it could lower. The tests are good on rank, axis and dtype. They miss several dynamic-shape and decline cases.

Code Quality

• exec_put_along_axis does not follow the MLX op's contract (runtime/MLXInterpreter.h:904-933). The other executors in this area are 1:1 wrappers: exec_take_along_axis (:185) and exec_scatter_add (:894). This one compares indices.shape(d) against x.shape(d) at runtime and does slice → put_along_axis → slice_update to match aten's leading-block rule.

  • Where the logic lives: the aten-to-MLX translation belongs in the Python handler, which is where this PR already puts it for gather (SliceNode on self) and for scatter's src.
  • Split responsibility: scatter's narrowing is now split. src is narrowed in Python and self in C++. A reader of PutAlongAxisNode can't tell from the schema that the node means "scatter into the leading index-shaped block". Anyone else who emits the node (a pattern, or another handler) gets aten semantics without asking for them.
  • Suggested fix: _scatter_handler can narrow x per axis with SliceNode, emit PutAlongAxisNode on the block, then write it back with SliceUpdateNode in reverse axis order. Both nodes already take IntOrVid start/stop. That leaves exec_put_along_axis as a single put_along_axis(...) call, the same shape as exec_scatter_add.
  • Fix this →

• Dynamic scalar value is declined instead of lowered (ops.py:2760). The handler picks the src path with isinstance(src, Slot). A scatter.value whose value is symbolic (for example from .item()) also arrives as a Slot, a Vid slot rather than a Tid. That sends it down the tensor branch, where n.args[3].meta["val"].shape raises AttributeError on a SymFloat/SymInt, and the node falls back to CPU. emit_lifted_constant already handles dynamic Slots by emitting FullNode(v=FloatOrVid). Branch on n.target == torch.ops.aten.scatter.value (or check src.id_type) so dynamic values take that path.

API Design

• _slice_to_shape rejects symbolic sizes it could lower with a Vid (ops.py:2653, :2665). The docstring says a symbolic size "has to be the same symbol on both sides, since the slice stop is fixed at export time". SliceNode.stop is IntOrVid, and emit_shape / SymSizeNode already exist to read a dim at runtime. Two cases are declined today and would work if the handler emitted SymSizeNode(index, axis) and passed the resulting Vid as stop:

  • index.size(d) is a different symbol from self.size(d) (for example s1 <= s0).
  • index.size(d) is symbolic while self.size(d) is static.

  This also covers the scatter-on-self narrowing above once it moves to Python, and it removes the only part of the current C++ design that works better (runtime shape comparison). The str(have) == str(want) check can stay as a fast path that skips the slice.

  The schema needs no other dynamism changes:

  • axis: int32 matches TakeAlongAxisNode/ScatterAddNode, and aten's dim is never symbolic.
  • values: Tid is enough, because scalars are lifted to 0-D tensors.

Testing

The rank, axis, negative-dim, dtype, transposed-view and exact-tolerance coverage is good. These gaps remain:

• Size-1 index on a non-scatter axis for scatter. This case matters most, because without narrowing MLX would broadcast and write every row. Gather covers it (index_shape=(1, 6, 3)). The scatter config index_shape=(1, 3, 5) at test_ops.py:7630 puts the 1 on the scatter axis (dim=0), so the non-scatter sizes are 3 and 5. Add, for example, shape=(4, 8), dim=1, index_shape=(1, 3).
• Dynamic dim on the gather/scatter axis itself. Every dynamic config uses batch dim 0 with dim=1. Add a case where self is dynamic along dim, such as a KV-length-style x: {1: seq} with a static index, to exercise the skip_axis path under symbolic shapes.
• Dynamic src larger than a static index. In the second scatter dynamic config, src gets index_shape, so the src-narrowing slice is never emitted for a symbolic src.
• Different symbolic dims on self and index (Dim("a") vs Dim("b")). Add one test either way: it should lower once the Vid change lands, or show it declines cleanly if it doesn't.
• Dynamic scatter.value (value from .item()), as described above.
• Decline paths are untested. test_partitioner.py already has tests that assert an op declines (for example test_zero_head_count_declines_instead_of_raising). Add tests that int64/float64 self for scatter, sparse_grad=True, and 0-D inputs stay on CPU and still give correct results end to end. The 8-byte decline is a GPU-kernel limitation, so it would regress silently if an MLX update added the kernel or if the check were removed.
• Lower priority: a zero-size index along dim for both ops, and a bool scatter. Masks are a common scatter target, and bool gets past the itemsize == 8 check.

Recommendation

Needs Discussion. The main point is the design question @metascroy raised: keep PutAlongAxisNode/exec_put_along_axis a direct mlx::put_along_axis wrapper and move the leading-block handling into the handler with IntOrVid stops. With that and the Vid change, the handler would decline fewer dynamic graphs. The other items are small fixes and extra tests.

[metascroy]

compare to int8/unint8 directly?

[metascroy]

Could this move to op_helpers.py?

[metascroy]

Thanks for the PR! The scatter lowering looks right to me. Narrowing self and src to the index block and then writing back with SliceUpdateNode handles the MLX broadcast mismatch, and the tests cover it well. Two things to fix before landing are inline.

[metascroy]

Please check x_meta.dtype.itemsize == 8 here, the same way _scatter_add_handler does. ScatterAxis::eval_gpu rejects every 8-byte dtype, not only int64 (mlx/backend/metal/indexing.cpp), so matching int64 exactly depends on the delegate never receiving float64 or another 8-byte type, and that could change.

[metascroy]

emit_lifted_constant caches constants under f"_scalar_{value}", and the dtype isn't part of the key. With the new float-into-int32 case (value=2.5, dtype=torch.int32), this lifts torch.tensor(2.5, dtype=torch.int32), which is 2, and caches it as _scalar_2.5. If a later op in the same graph lifts a float 2.5 (say y * 2.5), it gets the cached int tensor back and silently computes with 2.

The problem is in the helper, but this is the first handler that intentionally lifts a float value into an int dtype. Two ways to fix it:

• include the dtype in the cache key in emit_lifted_constant, or
• lift the scalar in its own Python type here, since mlx::put_along_axis already astypes values to a.dtype().

A test with a second float lift of the same value in the graph would cover this.