(sharktank) Python3.12+ not supported for torch.compile

[IanNod]

When getting setup with shortfin I see several sharktank test failures due to the Python3.12 version shortfin requires.

To recreate following setup instructions on https://github.com/nod-ai/SHARK-Platform/tree/main?tab=readme-ov-file#development-getting-started with python3.12

pytest sharktank/tests/ops/ops_test.py

Errors with

>       ep = torch.export.export(my_module, (a, d, qs, m))

sharktank/tests/ops/ops_test.py:289:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
venv/lib/python3.12/site-packages/torch/export/__init__.py:174: in export
    return _export(
venv/lib/python3.12/site-packages/torch/export/_trace.py:635: in wrapper
    raise e
venv/lib/python3.12/site-packages/torch/export/_trace.py:618: in wrapper
    ep = fn(*args, **kwargs)
venv/lib/python3.12/site-packages/torch/export/exported_program.py:83: in wrapper
    return fn(*args, **kwargs)
venv/lib/python3.12/site-packages/torch/export/_trace.py:860: in _export
    gm_torch_level = _export_to_torch_ir(
venv/lib/python3.12/site-packages/torch/export/_trace.py:347: in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py:1202: in inner
    check_if_dynamo_supported()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

    def check_if_dynamo_supported():
        if sys.version_info >= (3, 12):
>           raise RuntimeError("Python 3.12+ not yet supported for torch.compile")
E           RuntimeError: Python 3.12+ not yet supported for torch.compile

venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py:593: RuntimeError

[stellaraccident]

We should relax that to 3.10 to 3.13 and at least add unit test coverage for the older versions.

We'll want to use the newer versions for real runs because it is faster.

[marbre]

torch.compile is supported for Python 3.12 since torch 2.4.0, see https://dev-discuss.pytorch.org/t/torch-compile-support-for-python-3-12-completed/2054. We explicitly pin for the CPU
https://github.com/nod-ai/SHARK-Platform/blob/c6c7321e35aae8e3385d8232b51074e58becd053/pytorch-cpu-requirements.txt#L3
whereas we do not pin for the GPU
https://github.com/nod-ai/SHARK-Platform/blob/c6c7321e35aae8e3385d8232b51074e58becd053/pytorch-rocm-requirements.txt#L3

[marbre]

shortfin is now also usable and tested with Python 3.11 but tests hang for Python 3.10 (therefore it is not enabled in the CI).

[IanNod]

I'm assuming we are pinning torch to 2.3 due to previous issues with iree-turbine. That should be resolved now so I think we can bump torch without hopefully too much issue now.

[marbre]

I'm assuming we are pinning torch to 2.3 due to previous issues with iree-turbine. That should be resolved now so I think we can bump torch without hopefully too much issue now.

In the CI workflow iree-turbine's unit tests recently passed with torch==2.5.1 (see the logs here). Within a quick test, I bumped torch to 2.4.1 or 2.5.1 as part of #372 but this failed (see the logs here).

[marbre]

I've filled #495, with regards to test failures that occur for newer torch versions.