fix: preserve multi-value shape tensor inputs

py/torch_tensorrt/dynamo/utils.py

@@ -184,12 +184,13 @@ def get_torch_tensor(
     input: Input,
     device: torch.device,
     mode: str = "",
-) -> Union[int, torch.Tensor]:
+) -> Union[int, List[int], torch.Tensor]:
     if input.is_shape_tensor:
-        # TODO: All the shape tensors we've encountered so far are plain integers.
-        # Validate this assumption on more models.
         assert isinstance(input.shape, dict)
-        return input.shape["opt_shape"][0]
+        opt_shape = input.shape["opt_shape"]
+        # Scalar shape tensors supply SymInt parameters, while multi-element
+        # shape tensors supply SymInt[] parameters such as aten.full's size.
+        return opt_shape[0] if len(opt_shape) == 1 else list(opt_shape)
 
     if len(mode) > 0:
         return input.example_tensor(mode).to(device)
@@ -201,7 +202,10 @@ def get_torch_inputs(
     inputs: Sequence[Input] | Dict[str, Any],
     device: Union[Device, torch.device, str],
     mode: str = "",
-) -> Sequence[Union[int, torch.Tensor]] | Dict[str, Union[int, torch.Tensor]]:
+) -> (
+    Sequence[Union[int, List[int], torch.Tensor]]
+    | Dict[str, Union[int, List[int], torch.Tensor]]
+):
     """
     Return the torch_tensor from the Input object. If mode is set, this implies
     user is using dynamic shaped inputs and return the corresponding input based
@@ -210,15 +214,15 @@ def get_torch_inputs(
     device = to_torch_device(device)
 
     if isinstance(inputs, dict):
-        result_dict: Dict[str, Union[int, torch.Tensor]] = {}
+        result_dict: Dict[str, Union[int, List[int], torch.Tensor]] = {}
         for k, v in inputs.items():
             if isinstance(v, (list, tuple, dict)):
                 result_dict[k] = get_torch_inputs(v, device)
             elif isinstance(v, Input):
                 result_dict[k] = get_torch_tensor(v, device, mode)
         return result_dict
     else:
-        result_list: List[Union[int, torch.Tensor]] = []
+        result_list: List[Union[int, List[int], torch.Tensor]] = []
         for input in inputs:
             if isinstance(input, Input):
                 result_list.append(get_torch_tensor(input, device, mode))

tests/py/dynamo/runtime/test_000_compiler_utils.py

@@ -3,6 +3,7 @@
 import torch
 import torch_tensorrt
 from torch_tensorrt.dynamo.utils import (
+    get_torch_tensor,
     prepare_inputs,
     to_torch_device,
     to_torch_tensorrt_device,
@@ -57,6 +58,21 @@ def test_cast_str_device(self):
         self.assertTrue(prepared_device.gpu_id == gpu_id)
 
 
+class TestGetTorchTensor(unittest.TestCase):
+    def test_shape_tensor_preserves_multiple_values(self):
+        shape = torch_tensorrt.Input(
+            min_shape=(3, 5),
+            opt_shape=(3, 7),
+            max_shape=(4, 10),
+            dtype=torch.int64,
+            is_shape_tensor=True,
+        )
+
+        value = get_torch_tensor(shape, torch.device("cpu"))
+
+        self.assertEqual(value, [3, 7])
+
+
 class TestPrepareInputs(unittest.TestCase):
     def test_prepare_single_tensor_input(self):
         inputs = [torch.ones((4, 4))]