addressing review comment- unifying the shape functionality for upsample with concat

Author: apbose

py/torch_tensorrt/dynamo/conversion/impl/cat.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Union
+from typing import List, Optional, Sequence, Union
 
 import numpy as np
 import tensorrt as trt
@@ -16,6 +16,63 @@
 )
 
 
+def unify_trt_tensors(
+    ctx: ConversionContext,
+    target: Target,
+    name: str,
+    inputs: Sequence[Union[int, np.ndarray, torch.Tensor, TRTTensor]],
+    concat_axis: int,
+    cast_dtype: Union[_enums.dtype, trt.DataType, np.dtype] = None,
+    force_trt_output: bool = False,
+) -> Union[TRTTensor, List[int]]:
+    """
+    Normalize all inputs to TRT tensors if needed, optionally cast, and concat if any dynamic.
+
+    Args:
+        ctx: TensorRT conversion context.
+        target: FX target for naming.
+        name: Base name for layers.
+        inputs: Sequence of ints / numpy arrays / torch tensors / TRT tensors.
+        concat_axis: Axis along which to concatenate tensors if dynamic.
+        cast_dtype: Optional target dtype for casting TRT tensors.
+        force_trt_output: If True, return TRT tensor even if all inputs are static ints.
+    """
+    has_dynamic = any(not isinstance(x, int) for x in inputs)
+    trt_tensors = []
+
+    for i, x in enumerate(inputs):
+        # convert to TRTTensor
+        if isinstance(x, TRTTensor):
+            t = x
+        elif isinstance(x, int) and not has_dynamic and not force_trt_output:
+            t = x  # pure static path
+        else:
+            t = ctx.net.add_constant((1,), np.array([x], dtype=np.int32))
+            set_layer_name(t, target, f"{name}_dim{i}_const")
+            t = t.get_output(0)
+
+        # optional cast
+        if cast_dtype and isinstance(t, TRTTensor):
+            t = cast_trt_tensor(ctx, t, cast_dtype, f"{name}_cast_{i}")
+
+        trt_tensors.append(t)
+
+    if not has_dynamic and not force_trt_output:
+        return trt_tensors  # all ints
+
+    # promote remaining ints to TRT consts before concat
+    for i, t in enumerate(trt_tensors):
+        if isinstance(t, int):
+            const = ctx.net.add_constant((1,), np.array([t], dtype=np.int32))
+            set_layer_name(const, target, f"{name}_static_{i}_const")
+            trt_tensors[i] = const.get_output(0)
+
+    concat = ctx.net.add_concatenation(trt_tensors)
+    concat.axis = concat_axis
+    set_layer_name(concat, target, f"{name}_concat")
+    return concat.get_output(0)
+
+
 def cat(
     ctx: ConversionContext,
     target: Target,
@@ -54,9 +111,16 @@ def cat(
             )
             trt_casted_inputs.append(casted_input)
         trt_inputs = trt_casted_inputs
+    else:
+        trt_promoted_type = None
 
-    concat_layer = ctx.net.add_concatenation(trt_inputs)
     dim = get_positive_dim(dim, len(trt_inputs[0].shape))
-    concat_layer.axis = dim
-    set_layer_name(concat_layer, target, f"{name}_gather", source_ir)
-    return concat_layer.get_output(0)
+    return unify_trt_tensors(
+        ctx,
+        target,
+        name,
+        trt_inputs,
+        concat_axis=dim,
+        cast_dtype=trt_promoted_type,
+        force_trt_output=True,
+    )

py/torch_tensorrt/dynamo/conversion/impl/shape.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List, Optional, Tuple
+from typing import List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import tensorrt as trt
@@ -159,3 +159,52 @@ def to_trt_shape_tensor(
 
     # If no ITensor found, return plain list of ints
     return shape_list
+
+
+def collect_and_concat_trt_inputs(
+    ctx: ConversionContext,
+    target: Target,
+    name: str,
+    inputs: Sequence[Union[int, TRTTensor, torch.Tensor, np.ndarray]],
+    concat_axis: int = 0,
+    allow_static_return: bool = False,
+) -> Union[TRTTensor, List[int]]:
+    """
+    Normalize a sequence of values into TRT ITensors and concatenate them.
+    If `allow_static_return=True` and all inputs are ints, return a Python
+    list of ints instead of creating any TRT layers.
+    """
+    trt_tensors = []
+    has_dynamic = False
+
+    for i, x in enumerate(inputs):
+        if isinstance(x, TRTTensor):
+            trt_tensors.append(x)
+            has_dynamic = True
+
+        elif isinstance(x, (int, np.integer)):
+            # keep raw for now, convert only if dynamic found
+            trt_tensors.append(int(x))
+
+        else:
+            # torch/np tensor -> TRT tensor
+            t = get_trt_tensor(ctx, x, f"{name}_tensor_{i}")
+            trt_tensors.append(t)
+            has_dynamic = True
+
+    # fully static shape case
+    if not has_dynamic and allow_static_return:
+        return [int(v) for v in trt_tensors]
+
+    # promote remaining ints to TRT constants
+    for i, v in enumerate(trt_tensors):
+        if isinstance(v, int):
+            const = ctx.net.add_constant((1,), np.array([v], dtype=np.int32))
+            set_layer_name(const, target, f"{name}_static_dim{i}_const")
+            trt_tensors[i] = const.get_output(0)
+
+    # concatenate
+    concat = ctx.net.add_concatenation(trt_tensors)
+    concat.axis = concat_axis
+    set_layer_name(concat, target, f"{name}_concat")
+    return concat.get_output(0)

py/torch_tensorrt/dynamo/conversion/impl/upsample.py

@@ -9,9 +9,11 @@
     has_dynamic_shape,
     set_layer_name,
 )
+from torch_tensorrt.dynamo.conversion.impl.cat import (
+    unify_trt_tensors as unify_trt_shape_tensors,
+)
 from torch_tensorrt.dynamo.conversion.impl.shape import (
     get_shape_with_dynamic_shape,
-    to_trt_shape_tensor,
 )
 
 
@@ -40,7 +42,9 @@ def upsample(
             )
             layer.set_input(1, shape)
         else:
-            trt_shape = to_trt_shape_tensor(ctx, target, name, shape)
+            trt_shape = unify_trt_shape_tensors(
+                ctx, target, name, shape, concat_axis=0, force_trt_output=False
+            )
             if isinstance(trt_shape, list):
                 layer.shape = trt_shape
             else: