From 2f46d12c5c7a2b41a598962acca013c8944e35e2 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Wed, 9 Sep 2026 21:03:51 -0700 Subject: [PATCH] fix: support dynamic shapes in gather ## Problem `gather` selects values using an index tensor. It falls back to PyTorch whenever an input dimension can change, even though its TensorRT converter does not need fixed dimension sizes. Enabling that path without checks would also admit inputs the engine cannot handle safely. ## Change Declare dynamic-shape support and keep unsupported cases in PyTorch: an empty index, uint8 data, and float64 data unless `truncate_double=True` allows float32 arithmetic. An empty output can prevent the engine from running, affecting its other outputs too. Use Dynamo tracing in the dynamic tests so they actually exercise the support flag. Use nonzero values and varying indices to check the selected axis. ## Tests Passed 11/11 gather tests. The three dynamic conversion cases fail without this change. They cover positive and negative axes and a gathered dimension whose size can vary. Separate checks ran batch sizes 1, 3, and 6. Supported float32 gather built an engine and matched PyTorch. Float64, uint8, and empty-index cases fell back correctly. The empty-index check also verified a second output that ran in TensorRT. The standard dynamic harness executes the maximum configured shape; the separate batch checks cover the other sizes. A dynamic range containing zero was not tested. Tested on Linux x86_64 with Python 3.12, TensorRT 11.2, and the native runtime. Windows, aarch64, TensorRT-RTX, the Python-only runtime, and multiple GPUs were not tested. Existing out-of-range index behavior and conversion of indices to 32-bit integers are unchanged. --- .../dynamo/conversion/aten_ops_converters.py | 42 +++++++++++- .../py/dynamo/conversion/test_gather_aten.py | 67 +++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 898cd78ed1..c273ef12b2 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1404,7 +1404,47 @@ def aten_ops_clamp( ) -@dynamo_tensorrt_converter(torch.ops.aten.gather.default) +def gather_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Keep cases the TensorRT gather cannot serve on the PyTorch path. + + An empty index gives the engine a zero size output binding, which TensorRT refuses to + enqueue. The call still returns, so every other output of that engine silently comes + back zeroed. An engine built for a float64 input expects float32 and rejects the + caller's tensor at run time unless truncate_double is set, and a uint8 output fails the + engine build outright. All three ran correctly in PyTorch before this converter + accepted dynamic shapes. + """ + data_meta = ( + node.args[0].meta.get("tensor_meta") if hasattr(node.args[0], "meta") else None + ) + index_meta = ( + node.args[2].meta.get("tensor_meta") if hasattr(node.args[2], "meta") else None + ) + if index_meta is not None and 0 in tuple(index_meta.shape): + _LOGGER.debug("gather with an empty index is not supported, falling back") + return False + if data_meta is None: + return True + if data_meta.dtype == torch.uint8: + _LOGGER.debug("gather with a uint8 input is not supported, falling back") + return False + if data_meta.dtype == torch.float64 and not ( + settings is not None and settings.truncate_double + ): + _LOGGER.debug( + "gather with a float64 input needs truncate_double=True, falling back" + ) + return False + return True + + +@dynamo_tensorrt_converter( + torch.ops.aten.gather.default, + capability_validator=gather_validator, + supports_dynamic_shapes=True, +) @enforce_tensor_types( { 0: (TRTTensor,), diff --git a/tests/py/dynamo/conversion/test_gather_aten.py b/tests/py/dynamo/conversion/test_gather_aten.py index b4cbe475b4..926b142ec7 100644 --- a/tests/py/dynamo/conversion/test_gather_aten.py +++ b/tests/py/dynamo/conversion/test_gather_aten.py @@ -73,3 +73,70 @@ def forward(self, input, index): input = torch.zeros(3, 5, dtype=torch.int32) inputs = [input, index] self.run_test(TestModule(), inputs) + + @parameterized.expand( + [ + ("positive_dim", 1), + ("negative_dim", -1), + ] + ) + def test_gather_dynamic_shape(self, _, dim): + """The registry only consults supports_dynamic_shapes when a node carries symbolic + shape metadata, which the legacy tracer does not produce, so this needs the dynamo + tracer or it passes either way.""" + + class TestModule(torch.nn.Module): + def forward(self, input, index): + return torch.ops.aten.gather.default(input, dim, index) + + input_specs = [ + Input( + min_shape=(1, 5), + opt_shape=(3, 5), + max_shape=(6, 5), + dtype=torch.float32, + ), + Input( + min_shape=(1, 4), + opt_shape=(3, 4), + max_shape=(6, 4), + dtype=torch.int64, + ), + ] + self.run_test_with_dynamic_shape( + TestModule(), + input_specs, + use_dynamo_tracer=True, + ) + + def test_gather_dynamic_gathered_axis(self): + """The gathered axis itself is dynamic here, so index validity depends on the shape + the engine is given at run time rather than on the shape it was built at.""" + + class TestModule(torch.nn.Module): + def forward(self, input, index): + return torch.ops.aten.gather.default(input, 1, index) + + input_specs = [ + Input( + min_shape=(3, 1), + opt_shape=(3, 4), + max_shape=(3, 6), + dtype=torch.float32, + ), + Input( + min_shape=(3, 1), + opt_shape=(3, 4), + max_shape=(3, 6), + dtype=torch.int64, + ), + ] + self.run_test_with_dynamic_shape( + TestModule(), + input_specs, + use_dynamo_tracer=True, + ) + + +if __name__ == "__main__": + run_tests()