torch.ops.aten.rand_like ONNX export failure

[kyunggeunlee]

Follow-up of #2993

Reproducible Example

import torch
import pytest


@pytest.mark.parametrize(
    "op", [
        torch.ops.aten.rand_like,
        torch.ops.aten.randint_like,
        torch.ops.aten.randn_like,
    ]
)
def test_x_like(op: torch._ops.OpOverloadPacket):
    class Model(torch.nn.Module):
        def forward(self, x):
            return op(x, memory_format=torch.preserve_format)

    torch.onnx.export(
        Model(),
        (torch.randn(10, 10),),
        "{op.__name__}.onnx",
        input_names=["input"],
        output_names=["output"],
        dynamo=True,
    )


if __name__ == "__main__":
    pytest.main([__file__])

[fw7th]

Don't understand what's going on yet, but I'm working on reproducing it locally. I'll update this thread once I have a lead.

[promiseeuler]

I reproduced this on current ONNXScript source with PyTorch 2.10.0 and ONNXScript 0.7.2.dev20260822. The first actionable exception is:

TypeError: aten_rand_like() got an unexpected keyword argument 'memory_format'

The same mismatch exists for all three functions in function_libs/torch_lib/ops/core.py: their docstrings include MemoryFormat? memory_format=None, but the Python signatures stop at pin_memory.

I tested a minimal local change that adds an ignored memory_format: str = "" parameter to aten_rand_like, aten_randint_like, and aten_randn_like. With that change, exports using memory_format=torch.preserve_format passed for all three:

rand_like PASS
randint_like PASS
randn_like PASS

This is consistent with the existing TorchLib treatment of memory format: ONNX tensors do not represent PyTorch memory layout, so functions such as aten_contiguous already treat this argument as a no-op.

One scope detail for the eventual fix: aten_randint_like_low_dtype has the same documented-but-missing parameter and should receive it too, even though the posted reproduction selects the default overload. A focused regression should retain the three parameterized exports from the issue and add the low/dtype overload so the signatures cannot drift again.

[fw7th]

I tested a minimal local change that adds an ignored memory_format: str = "" parameter to aten_rand_like, aten_randint_like, and aten_randn_like. With that change, exports using memory_format=torch.preserve_format passed for all three:

Thank you, I was confused about the function of memory_format considering it's never used in the function definitions in the aforementioned the aten_* functions, this really helps.

[justinchuby]

Please feel free to put up a PR. Thanks

[gabrielfruet]

7th (@fw7th) saw you were digging into this. the random *_like ops don't
accept the memory_format kwarg the schema declares, the same gap #2994 fixed
for full_like. Fix is up in #3011.