diff --git a/py/torch_tensorrt/dynamo/utils.py b/py/torch_tensorrt/dynamo/utils.py index 2dc4bdeece..38f01ed80e 100644 --- a/py/torch_tensorrt/dynamo/utils.py +++ b/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,7 +214,7 @@ 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) @@ -218,7 +222,7 @@ def get_torch_inputs( 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)) diff --git a/tests/py/dynamo/runtime/test_000_compiler_utils.py b/tests/py/dynamo/runtime/test_000_compiler_utils.py index 42d226911a..65f90883b6 100644 --- a/tests/py/dynamo/runtime/test_000_compiler_utils.py +++ b/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))]