(shartank) testExportWithArgumentDeviceAffinities fails with torch 2.4.0 and above

[marbre]

Using

• HEAD at commit d279aff
• Python 3.11.10
• iree-base-compiler==3.1.0rc20241212
• iree-base-runtime==3.1.0rc20241212
• iree-turbine==3.1.0rc20241211
• torch==2.5.1+cpu, torch==2.5.0+cpu, torch==2.4.1+cpu and torch==2.4.0+cpu

the test testExportWithArgumentDeviceAffinities fails with

FAILED tests/export_test.py::ExportTest::testExportWithArgumentDeviceAffinities - torch._dynamo.exc.UserError: Expecting args to be a tuple of example positional inputs, got <class 'list'>

Full error log:

self = <tests.export_test.ExportTest testMethod=testExportWithArgumentDeviceAffinities>

    def testExportWithArgumentDeviceAffinities(self):
        args = (ReplicatedTensor(ts=[torch.tensor([1])]), torch.tensor([[2]]))
    
        class Module(torch.nn.Module):
            def f(self, a, b):
                return a, b
    
        module = Module()
        fxb = FxProgramsBuilder(module)
>       export(
            Module.f,
            fx_builder=fxb,
            args=args,
            strict=False,
        )

tests/export_test.py:82: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
sharktank/export.py:166: in export
    return fx_builder.export_program(
.venv-uv-3.11/lib/python3.11/site-packages/iree/turbine/aot/fx_programs.py:239: in export_program
    program = torch.export.export(
.venv-uv-3.11/lib/python3.11/site-packages/torch/export/__init__.py:270: in export
    return _export(
.venv-uv-3.11/lib/python3.11/site-packages/torch/export/_trace.py:1017: in wrapper
    raise e
.venv-uv-3.11/lib/python3.11/site-packages/torch/export/_trace.py:990: in wrapper
    ep = fn(*args, **kwargs)
.venv-uv-3.11/lib/python3.11/site-packages/torch/export/exported_program.py:114: in wrapper
    return fn(*args, **kwargs)
.venv-uv-3.11/lib/python3.11/site-packages/torch/export/_trace.py:1875: in _export
    ) = _process_export_inputs(mod, args, kwargs, dynamic_shapes)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

mod = LambdaModule(
  (root): Module()
), args = [tensor([1]), tensor([[2]])], kwargs = None, dynamic_shapes = None

    def _process_export_inputs(mod, args, kwargs, dynamic_shapes):
        original_state_dict = mod.state_dict(keep_vars=True)
    
        if not isinstance(args, tuple):
>           raise UserError(
                UserErrorType.INVALID_INPUT,
                f"Expecting `args` to be a tuple of example positional inputs, got {type(args)}",
            )
E           torch._dynamo.exc.UserError: Expecting `args` to be a tuple of example positional inputs, got <class 'list'>

.venv-uv-3.11/lib/python3.11/site-packages/torch/export/_trace.py:1049: UserError