Function _embedding_handler

backends/mlx/ops.py:1076–1095

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.embedding.default])
def _embedding_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    args = P.args(n)
    require_args(args, 2, 3, "aten.embedding")
    # "padding_idx", "scale_grad_by_freq", "sparse" are training only args
    # and ignored
    require_kwargs(
        P.kwargs(n), {"padding_idx", "scale_grad_by_freq", "sparse"}, "aten.embedding"
    )
    w, x = args[0], args[1]
    # padding_idx (args[2] if present) is ignored - only affects gradients
    out = P.make_or_get_slot(n)
    P.emit(
        TakeNode(
            x=P.slot_to_tid(w),
            index=IntOrVidOrTid.from_tid(P.slot_to_tid(x)),
            out=P.slot_to_tid(out),
            axis=0,
        )
    )
    return out


@REGISTRY.register(target=[torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar])