(P: MLXProgramBuilder, n: Node)
| 1728 | |
| 1729 | @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) |
| 1730 | def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: |
| 1731 | args = P.args(n) |
| 1732 | require_args(args, 2, 2, "aten.index.Tensor") |
| 1733 | require_kwargs(P.kwargs(n), set(), "aten.index.Tensor") |
| 1734 | x, idx_list = args |
| 1735 | if not isinstance(idx_list, list) or len(idx_list) == 0: |
| 1736 | raise ValueError( |
| 1737 | f"aten.index.Tensor requires a list of index tensors, " |
| 1738 | f"got {type(idx_list)}" |
| 1739 | ) |
| 1740 | |
| 1741 | x_meta = n.args[0].meta.get("val") |
| 1742 | x_ndim = len(x_meta.shape) if x_meta is not None else None |
| 1743 | |
| 1744 | # Filter out None indices and track which axes they correspond to |
| 1745 | non_none = [(i, idx) for i, idx in enumerate(idx_list) if idx is not None] |
| 1746 | |
| 1747 | if len(non_none) == 0: |
| 1748 | raise ValueError("aten.index.Tensor: all indices are None") |
| 1749 | |
| 1750 | if len(non_none) == 1: |
| 1751 | axis, idx = non_none[0] |
| 1752 | idx_meta = n.args[1][axis].meta.get("val") |
| 1753 | ndim_match = ( |
| 1754 | x_meta is not None |
| 1755 | and idx_meta is not None |
| 1756 | and len(x_meta.shape) == len(idx_meta.shape) |
| 1757 | ) |
| 1758 | out = P.make_or_get_slot(n) |
| 1759 | if ndim_match: |
| 1760 | # Same ndim: use TakeAlongAxisNode (element-wise gather) |
| 1761 | P.emit( |
| 1762 | TakeAlongAxisNode( |
| 1763 | x=P.slot_to_tid(x), |
| 1764 | indices=P.slot_to_tid(idx), |
| 1765 | out=P.slot_to_tid(out), |
| 1766 | axis=axis, |
| 1767 | ) |
| 1768 | ) |
| 1769 | else: |
| 1770 | # Different ndim (e.g. 1D indices into 3D tensor): use TakeNode |
| 1771 | P.emit( |
| 1772 | TakeNode( |
| 1773 | x=P.slot_to_tid(x), |
| 1774 | index=IntOrVidOrTid.from_tid(P.slot_to_tid(idx)), |
| 1775 | out=P.slot_to_tid(out), |
| 1776 | axis=axis, |
| 1777 | ) |
| 1778 | ) |
| 1779 | return out |
| 1780 | |
| 1781 | # Multi-index: use GatherNode (maps to mlx::gather) |
| 1782 | if x_meta is None or x_ndim is None: |
| 1783 | raise ValueError( |
| 1784 | "aten.index.Tensor with multiple indices requires input shape metadata" |
| 1785 | ) |
| 1786 | |
| 1787 | indices = [P.slot_to_tid(idx) for _, idx in non_none] |
nothing calls this directly
no test coverage detected