Function _add_handler

backends/mlx/ops.py:1099–1129

Handle aten.add.Tensor: a + alpha * b.

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar])
def _add_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle aten.add.Tensor: a + alpha * b."""
    args = P.args(n)
    require_args(args, 2, 2, "aten.add.Tensor")
    require_kwargs(P.kwargs(n), {"alpha"}, "aten.add.Tensor")
    a, b = args
    input_meta = n.args[0].meta.get("val")
    dtype = input_meta.dtype if input_meta is not None else torch.float32
    if not isinstance(b, Slot):
        b = emit_lifted_constant(P, b, dtype)
    alpha = P.kwargs(n).get("alpha", 1)
    if alpha != 1:
        alpha_slot = emit_lifted_constant(P, alpha, dtype)
        _, tmp = P.make_tmp_slot()
        P.emit(
            MultiplyNode(
                a=P.slot_to_tid(b),
                b=P.slot_to_tid(alpha_slot),
                out=P.slot_to_tid(tmp),
            )
        )
        b = tmp
    out = P.make_or_get_slot(n)
    P.emit(
        AddNode(
            a=P.slot_to_tid(a),
            b=P.slot_to_tid(b),
            out=P.slot_to_tid(out),
        )
    )
    return out


@REGISTRY.register(target=[torch.ops.aten.div.Tensor_mode])