Function _linear_handler

backends/mlx/ops.py:790–816

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.linear.default])
def _linear_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    args = P.args(n)
    require_args(args, 2, 3, "aten.linear")
    require_kwargs(P.kwargs(n), set(), "aten.linear")
    x, w = args[0], args[1]
    b = args[2] if len(args) > 2 else None
    out = P.make_or_get_slot(n)

    # Transpose weight: linear(x, w) = x @ w.T
    _, w_t = P.make_tmp_slot()
    P.emit(
        TransposeNode(
            x=P.slot_to_tid(w),
            out=P.slot_to_tid(w_t),
            perm=[1, 0],
        )
    )

    P.emit(
        AddmmNode(
            mat1=P.slot_to_tid(x),
            mat2=P.slot_to_tid(w_t),
            out=P.slot_to_tid(out),
            bias=P.slot_to_tid(b) if b else None,
        )
    )
    return out


@REGISTRY.register(target=[torch.ops.aten.addmm.default])