fix: support dynamic slice_scatter fallback shapes

py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py

@@ -22,8 +22,11 @@
 from typing import Optional, Tuple
 
 import numpy as np
+import tensorrt as trt
+from tensorrt import ITensor as TRTTensor
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import (
     AliasedOutput,
     AliasKind,
@@ -34,9 +37,8 @@
     set_layer_name,
 )
 from torch_tensorrt.dynamo.conversion.impl import select
-
-import tensorrt as trt
-from tensorrt import ITensor as TRTTensor
+from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
+from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
 
 logger = logging.getLogger(__name__)
 
@@ -215,8 +217,36 @@ def slice_scatter(
     target_shape[dim] = len(indices_np)
     indices_np = indices_np.reshape(target_shape)
     src_shape = tuple(src.shape)
-    indices_np = np.broadcast_to(indices_np, src_shape).astype(np.int64)
-    indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices")
+    if DYNAMIC_DIM in src_shape:
+        indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices")
+        runtime_src_shape = []
+        for axis, size in enumerate(src_shape):
+            if axis == dim % rank:
+                runtime_src_shape.append(len(indices_np))
+            elif size == DYNAMIC_DIM:
+                runtime_src_shape.append(
+                    get_shape(
+                        ctx,
+                        target,
+                        source_ir,
+                        name + f"_fallback_src_shape_{axis}",
+                        src,
+                        axis,
+                    )
+                )
+            else:
+                runtime_src_shape.append(size)
+        indices_tensor = impl.slice.expand(
+            ctx,
+            target,
+            source_ir,
+            name + "_fallback_indices_expand",
+            indices_tensor,
+            tuple(runtime_src_shape),
+        )
+    else:
+        indices_np = np.broadcast_to(indices_np, src_shape).astype(np.int64)
+        indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices")
 
     return select.scatter(
         ctx,

tests/py/dynamo/conversion/test_slice_scatter_aten.py

@@ -21,6 +21,7 @@
 import torch
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -76,6 +77,24 @@ def test_fallback_step_two(self):
         update = torch.randn(2, 4, 8, 8)
         self.run_test(module, [cache, update])
 
+    def test_fallback_dynamic_shape(self):
+        module = _SliceScatterNotInputModule(2, 1, 60, step=3)
+        input_specs = [
+            Input(
+                min_shape=(2, 4, 64, 8),
+                opt_shape=(3, 4, 64, 12),
+                max_shape=(4, 4, 64, 16),
+                dtype=torch.float32,
+            ),
+            Input(
+                min_shape=(2, 4, 20, 8),
+                opt_shape=(3, 4, 20, 12),
+                max_shape=(4, 4, 20, 16),
+                dtype=torch.float32,
+            ),
+        ]
+        self.run_test_with_dynamic_shape(module, input_specs)
+
     def test_full_overwrite_is_identity(self):
         """When start=0, end=dim_size, step=1, the converter short-circuits
         and returns ``src`` directly. Wrap the returned tensor in a small op