export_paged_llm_v1.py crashes for --bs=1

[NoumanAmir657]

Steps to reproduce

python -m sharktank.examples.export_paged_llm_v1 --hf-dataset=open_llama_3b_v2_f16_gguf --output-mlir=/tmp/open_llama_3b_v2_f16.mlir --output-config=/tmp/open_llama_3b_v2_f16.json --bs=1

Exporting prefill_bs1
W1206 13:02:32.103000 516072 .venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py:650] Using None as a dynamic shape dimension is deprecated. Please use Dim.STATIC instead
W1206 13:02:32.103000 516072 .venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py:650] Using None as a dynamic shape dimension is deprecated. Please use Dim.STATIC instead
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/home/nouman-10x/shark-ai/sharktank/sharktank/examples/export_paged_llm_v1.py", line 358, in <module>
    main()
  File "/home/nouman-10x/shark-ai/sharktank/sharktank/examples/export_paged_llm_v1.py", line 339, in main
    generate_batch_prefill(bs)
  File "/home/nouman-10x/shark-ai/sharktank/sharktank/examples/export_paged_llm_v1.py", line 208, in generate_batch_prefill
    @fxb.export_program(
     ^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/iree/turbine/aot/fx_programs.py", line 239, in export_program
    program = torch.export.export(
              ^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/__init__.py", line 270, in export
    return _export(
           ^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/_trace.py", line 1017, in wrapper
    raise e
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/_trace.py", line 990, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/exported_program.py", line 114, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/_trace.py", line 1880, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/_trace.py", line 1643, in _non_strict_export
    ) = make_fake_inputs(
        ^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 139, in make_fake_inputs
    transformed_dynamic_shapes = _transform_shapes_for_default_dynamic(
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 911, in _transform_shapes_for_default_dynamic
    result = _tree_map_with_path(
             ^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 481, in _tree_map_with_path
    return tree_map_with_path(f, tree, *dynamic_shapes, is_leaf=is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/utils/_pytree.py", line 1608, in tree_map_with_path
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/utils/_pytree.py", line 803, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/utils/_pytree.py", line 1608, in <genexpr>
    return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves))
                              ^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 478, in f
    return func(path, t, *dynamic_shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 909, in transform_shape
    return transform_shapes(path, t, dynamic_shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nouman-10x/shark-ai/sharktank/.venv/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 880, in transform_shapes
    dim = shape[i]
          ~~~~~^^^
IndexError: list index out of range

Does not seem to fail if I pass --bs=1,2.
But the cache type it will use then would not be direct.