backends/mlx: reduce index handler lint complexity

Author: seyeong-han

backends/mlx/ops.py

@@ -1702,82 +1702,88 @@ def _index_gather_permutation(
     )
 
 
-@REGISTRY.register(target=[torch.ops.aten.index.Tensor])
-def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
-    args = P.args(n)
-    require_args(args, 2, 2, "aten.index.Tensor")
-    require_kwargs(P.kwargs(n), set(), "aten.index.Tensor")
-    x, idx_list = args
+def _non_none_index_tensors(idx_list: Any) -> List[Tuple[int, Slot]]:
     if not isinstance(idx_list, list) or len(idx_list) == 0:
         raise ValueError(
-            f"aten.index.Tensor requires a list of index tensors, "
-            f"got {type(idx_list)}"
+            f"aten.index.Tensor requires a list of index tensors, got {type(idx_list)}"
         )
 
-    x_meta = n.args[0].meta.get("val")
-    x_ndim = len(x_meta.shape) if x_meta is not None else None
-
-    # Filter out None indices and track which axes they correspond to
     non_none = [(i, idx) for i, idx in enumerate(idx_list) if idx is not None]
-
     if len(non_none) == 0:
         raise ValueError("aten.index.Tensor: all indices are None")
+    return non_none
 
-    if len(non_none) == 1:
-        axis, idx = non_none[0]
-        idx_meta = n.args[1][axis].meta.get("val")
-        ndim_match = (
-            x_meta is not None
-            and idx_meta is not None
-            and len(x_meta.shape) == len(idx_meta.shape)
+
+def _emit_single_index_handler(
+    P: MLXProgramBuilder,
+    n: Node,
+    x: Slot,
+    axis: int,
+    idx: Slot,
+    x_meta: Any,
+) -> Slot:
+    idx_meta = n.args[1][axis].meta.get("val")
+    ndim_match = (
+        x_meta is not None
+        and idx_meta is not None
+        and len(x_meta.shape) == len(idx_meta.shape)
+    )
+    out = P.make_or_get_slot(n)
+    if ndim_match:
+        # Same ndim: use TakeAlongAxisNode (element-wise gather)
+        P.emit(
+            TakeAlongAxisNode(
+                x=P.slot_to_tid(x),
+                indices=P.slot_to_tid(idx),
+                out=P.slot_to_tid(out),
+                axis=axis,
+            )
         )
-        out = P.make_or_get_slot(n)
-        if ndim_match:
-            # Same ndim: use TakeAlongAxisNode (element-wise gather)
-            P.emit(
-                TakeAlongAxisNode(
-                    x=P.slot_to_tid(x),
-                    indices=P.slot_to_tid(idx),
-                    out=P.slot_to_tid(out),
-                    axis=axis,
-                )
+    else:
+        # Different ndim (e.g. 1D indices into 3D tensor): use TakeNode
+        P.emit(
+            TakeNode(
+                x=P.slot_to_tid(x),
+                index=IntOrVidOrTid.from_tid(P.slot_to_tid(idx)),
+                out=P.slot_to_tid(out),
+                axis=axis,
             )
-        else:
-            # Different ndim (e.g. 1D indices into 3D tensor): use TakeNode
-            P.emit(
-                TakeNode(
-                    x=P.slot_to_tid(x),
-                    index=IntOrVidOrTid.from_tid(P.slot_to_tid(idx)),
-                    out=P.slot_to_tid(out),
-                    axis=axis,
-                )
+        )
+    return out
+
+
+def _index_slice_sizes(x_meta: Any, x_ndim: int, indexed_axes: Set[int]) -> List[int]:
+    slice_sizes = []
+    for dim in range(x_ndim):
+        if dim in indexed_axes:
+            slice_sizes.append(1)
+            continue
+
+        dim_size = x_meta.shape[dim]
+        if not isinstance(dim_size, int):
+            raise ValueError(
+                f"aten.index.Tensor: non-indexed dimension {dim} has dynamic size "
+                f"{dim_size}, which is not supported with multi-index gather"
             )
-        return out
+        slice_sizes.append(dim_size)
+    return slice_sizes
 
-    # Multi-index: use GatherNode (maps to mlx::gather)
-    if x_meta is None or x_ndim is None:
-        raise ValueError(
-            "aten.index.Tensor with multiple indices requires input shape metadata"
-        )
 
+def _emit_multi_index_handler(
+    P: MLXProgramBuilder,
+    n: Node,
+    x: Slot,
+    x_meta: Any,
+    x_ndim: int,
+    non_none: List[Tuple[int, Slot]],
+) -> Slot:
     indices = [P.slot_to_tid(idx) for _, idx in non_none]
     axes = [i for i, _ in non_none]
+    indexed_axes = set(axes)
 
     # slice_sizes: 1 for indexed axes, full dim size for non-indexed axes
     # Use int() to handle SymInt values from dynamic shapes
-    indexed_axes = set(axes)
-    slice_sizes = []
-    for dim in range(x_ndim):
-        if dim in indexed_axes:
-            slice_sizes.append(1)
-        else:
-            dim_size = x_meta.shape[dim]
-            if not isinstance(dim_size, int):
-                raise ValueError(
-                    f"aten.index.Tensor: non-indexed dimension {dim} has dynamic size "
-                    f"{dim_size}, which is not supported with multi-index gather"
-                )
-            slice_sizes.append(dim_size)
+    slice_sizes = _index_slice_sizes(x_meta, x_ndim, indexed_axes)
 
     # Emit gather — output shape is broadcast(indices).shape + slice_sizes
     _, gather_slot = P.make_tmp_slot()
@@ -1841,6 +1847,29 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+@REGISTRY.register(target=[torch.ops.aten.index.Tensor])
+def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    args = P.args(n)
+    require_args(args, 2, 2, "aten.index.Tensor")
+    require_kwargs(P.kwargs(n), set(), "aten.index.Tensor")
+    x, idx_list = args
+
+    x_meta = n.args[0].meta.get("val")
+    x_ndim = len(x_meta.shape) if x_meta is not None else None
+    non_none = _non_none_index_tensors(idx_list)
+
+    if len(non_none) == 1:
+        axis, idx = non_none[0]
+        return _emit_single_index_handler(P, n, x, axis, idx, x_meta)
+
+    if x_meta is None or x_ndim is None:
+        raise ValueError(
+            "aten.index.Tensor with multiple indices requires input shape metadata"
+        )
+
+    return _emit_multi_index_handler(P, n, x, x_meta, x_ndim, non_none)
+
+
 @REGISTRY.register(target=[torch.ops.aten.index_select.default])
 def _index_select_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle aten.index_select: select elements along an axis using a 1D index tensor.