Skip to content

Commit 43cb2b2

Browse files
fix(mlx): forward boundary-partitioned getitem instead of indexing a Slot (#21441)
## Summary The MLX partitioner can place an `operator.getitem` inside a partition while leaving its multi-output parent (e.g. `aten.native_layer_norm`) outside it. In that case `_getitem_handler` receives the already-selected element as a single boundary-input `Slot` rather than the source tuple, so `a[idx]` raises `'Slot' object is not subscriptable` and aborts the export. This fix forwards the `Slot` directly when the resolved arg isn't a tuple/list — the partitioner has already fed the selected element as the partition input. ## Impact Unblocks transformer / layer-norm-heavy models (e.g. RF-DETR) through the MLX delegate. Localized, one-file change to the op handler; no schema/serialization change. ## Test plan - Verified that DETR-style exports (dozens of `native_layer_norm` ops) lower and run correctly on device through the MLX delegate with this change. - Non-boundary getitem (source is a tuple/list of slots) is unchanged: `a[idx]` still selects the element. cc @metascroy Co-authored-by: Scott Roy <161522778+metascroy@users.noreply.github.com>
1 parent 6898c99 commit 43cb2b2

1 file changed

Lines changed: 5 additions & 1 deletion

File tree

backends/mlx/ops.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2272,9 +2272,13 @@ def _getitem_handler(P: MLXProgramBuilder, n: Node) -> Slot:
22722272
require_kwargs(P.kwargs(n), set(), "operator.getitem")
22732273
a, idx = args
22742274
out = P.make_or_get_slot(n)
2275+
# When the multi-output parent is left outside the partition, the partitioner
2276+
# feeds this getitem the already-selected element as a single boundary-input
2277+
# Slot (not the source tuple). In that case forward it directly.
2278+
src = a[idx] if isinstance(a, (tuple, list)) else a
22752279
P.emit(
22762280
IdCopyNode(
2277-
x=P.slot_to_tid(a[idx]),
2281+
x=P.slot_to_tid(src),
22782282
out=P.slot_to_tid(out),
22792283
)
22802284
)

0 commit comments

Comments
 (0)