Function _rope_handler

backends/mlx/ops.py:2181–2213

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.mlx.rope.default])
def _rope_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    args = P.args(n)
    require_args(args, 3, 7, "mlx.rope")
    require_kwargs(P.kwargs(n), set(), "mlx.rope")
    x, dims, pos = args[0], args[1], args[2]
    traditional = args[3] if len(args) > 3 else False
    base = args[4] if len(args) > 4 else 500000.0
    scale = args[5] if len(args) > 5 else 1.0
    freqs = args[6] if len(args) > 6 else None
    out = P.make_or_get_slot(n)

    # pos must be a Slot (SymInt) from input_pos.item() during tracing
    # The schema supports both Vid (scalar) and Tid (tensor) for offset
    if not isinstance(pos, Slot):
        raise ValueError(
            f"RopeNode.offset must be a SymInt (traced via tensor.item()), got {type(pos)}. "
            "Make sure input_pos is a tensor and you call input_pos.item() to get a SymInt."
        )

    P.emit(
        RopeNode(
            x=P.slot_to_tid(x),
            out=P.slot_to_tid(out),
            dims=dims,
            offset=VidOrTid.from_vid(P.slot_to_vid(pos)),
            freqs=P.slot_to_tid(freqs) if freqs else None,
            traditional=traditional,
            base=base,
            scale=scale,
        )
    )

    return out


def _emit_channel_last_weight(P: MLXProgramBuilder, w_node: Node, perm: list) -> Slot: