[sharktank] Test failures with torch>=2.4.0

[marbre]

Using Python 3.11.9, sharktank at HEAD fails when using a more recent torch version. In the CI, torch is pinned to 2.3.0:

SHARK-Platform/.github/workflows/ci-sharktank.yml

Line 61 in 903d3c1

pip install --no-compile -r pytorch-cpu-requirements.txt

SHARK-Platform/pytorch-cpu-requirements.txt

Line 3 in 903d3c1

torch==2.3.0

When upgrading to a more recent version, the following tests fail:

torch==2.4.0 and torch==2.4.1

FAILED sharktank/tests/export_test.py::ExportTest::testExportWithArgumentDeviceAffinities - torch._dynamo.exc.UserError: Expecting args to be a tuple of example positional inputs, got <class 'list'>
FAILED sharktank/tests/layers/paged_llama_attention_block_test.py::PagedLlamaAttentionBlockTest::testExportNondecomposed - iree.compiler._mlir_libs._site_initialize..MLIRError: Verification failed:

torch==2.5.0 and torch==2.5.1

FAILED sharktank/tests/export_test.py::ExportTest::testExportWithArgumentDeviceAffinities - torch._dynamo.exc.UserError: Expecting args to be a tuple of example positional inputs, got <class 'list'>
FAILED sharktank/tests/layers/paged_llama_attention_block_test.py::PagedLlamaAttentionBlockTest::testExportNondecomposed - AssertionError: 'torch.aten._scaled_dot_product_flash_attention_for_cpu' not found in 'module @module {\n util.global private @__auto.constant_660_660_torch.float32 = dense_resource<__a...
FAILED sharktank/tests/layers/sharded_conv2d_with_iree_test.py::test_sharded_conv2d_with_iree - ValueError: Node type mismatch; expected <class 'tuple'>, but got <class 'list'>.
FAILED sharktank/tests/models/punet/sharded_resnet_block_with_iree_test.py::test_sharded_resnet_block_with_iree - ValueError: Node type mismatch; expected <class 'tuple'>, but got <class 'list'>.

[marbre]

torch<=2.4 also blocks us from unpinning numpy (we're currently restricted to a version <2.0.0).