Function _scalar_tensor_handler

backends/mlx/ops.py:3155–3190

This is equivalent to torch.full([], scalar, dtype=dtype).

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.scalar_tensor.default])
def _scalar_tensor_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """This is equivalent to torch.full([], scalar, dtype=dtype)."""
    args = P.args(n)
    kwargs = P.kwargs(n)
    require_args(args, 1, 1, "aten.scalar_tensor")
    require_kwargs(
        kwargs, {"dtype", "layout", "device", "pin_memory"}, "aten.scalar_tensor"
    )
    require_contiguous_format(
        layout=kwargs.get("layout"),
        op_name="aten.scalar_tensor",
    )
    scalar_value = args[0]

    out = P.make_or_get_slot(n)

    # Get dtype from kwargs, default to float32
    dtype = n.kwargs.get("dtype")
    if dtype is None:
        # Infer dtype from scalar type
        if isinstance(scalar_value, bool):
            dtype = torch.bool
        elif isinstance(scalar_value, int):
            dtype = torch.int64
        else:
            dtype = torch.float32

    P.emit(
        FullNode(
            out=P.slot_to_tid(out),
            shape=[],  # 0-D tensor (scalar)
            v=P.to_float_or_vid(scalar_value),
            scalar_type=torch_dtype_to_scalar_type(dtype),
        )
    )
    return out


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