(P: MLXProgramBuilder, n: Node)
| 1689 | |
| 1690 | @REGISTRY.register(target=[torch.ops.aten.roll.default]) |
| 1691 | def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot: |
| 1692 | args = P.args(n) |
| 1693 | require_args(args, 2, 3, "aten.roll") |
| 1694 | require_kwargs(P.kwargs(n), set(), "aten.roll") |
| 1695 | x = args[0] |
| 1696 | shifts_arg = args[1] |
| 1697 | dims_arg = args[2] if len(args) > 2 else [] |
| 1698 | |
| 1699 | shifts = [shifts_arg] if isinstance(shifts_arg, int) else list(shifts_arg) |
| 1700 | dims: List[int] = [dims_arg] if isinstance(dims_arg, int) else list(dims_arg) |
| 1701 | |
| 1702 | # Flat roll (torch.roll with dims=[]) would require reshape + roll + |
| 1703 | # reshape at the graph level. Not yet supported; Swin-style usage always |
| 1704 | # passes explicit dims. |
| 1705 | if not dims: |
| 1706 | raise NotImplementedError( |
| 1707 | "aten.roll without dims (flat roll) is not supported by the MLX " |
| 1708 | "delegate yet." |
| 1709 | ) |
| 1710 | if len(shifts) != len(dims): |
| 1711 | raise ValueError( |
| 1712 | f"aten.roll: shifts and dims must have the same length, got " |
| 1713 | f"shifts={shifts} (len={len(shifts)}) dims={dims} (len={len(dims)})" |
| 1714 | ) |
| 1715 | require_static_ints(dims, "dims", "aten.roll") |
| 1716 | |
| 1717 | out = P.make_or_get_slot(n) |
| 1718 | P.emit( |
| 1719 | RollNode( |
| 1720 | x=P.slot_to_tid(x), |
| 1721 | out=P.slot_to_tid(out), |
| 1722 | shift=[P.to_int_or_vid(s) for s in shifts], |
| 1723 | axes=dims, |
| 1724 | ) |
| 1725 | ) |
| 1726 | return out |
| 1727 | |
| 1728 | |
| 1729 | @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) |
nothing calls this directly
no test coverage detected