Add normalization of cat operator argument into valid input (#3890)

Author: apbose

examples/distributed_inference/data_parallel_stable_diffusion.py

@@ -53,7 +53,5 @@
 
 # Assume there are 2 processes (2 devices)
 with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
-    print("before \n")
     result = pipe(prompt).images[0]
-    print("after ")
     result.save(f"result_{distributed_state.process_index}.png")

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2,7 +2,7 @@
 
 import logging
 import operator
-from typing import Callable, Dict, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -218,21 +218,57 @@ def aten_ops_native_group_norm(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.cat.default, supports_dynamic_shapes=True)
+def parse_cat_args(
+    args: Tuple[Argument, ...], kwargs: Dict[str, Any]
+) -> Tuple[List[Any], int]:
+    """
+    Process inputs for torch.ops.aten.cat.default.
+
+    Handles these valid patterns:
+      1. args = ((t1, t2, ...), dim)
+      2. args = ((t1, t2, ...),), kwargs = {dim: X} with optional dim in kwargs
+
+    Returns:
+        (input_tensors, dim)
+        input_tensors: tuple of tensor arguments
+        dim: integer concatenation dimension (default 0)
+    """
+
+    if len(args) > 1 and isinstance(args[0], (list, tuple)):
+        input_tensors = list(args[0])
+        dim = args_bounds_check(args, 1, 0)
+
+    else:
+        # If single arg is itself a tuple/list, unwrap it
+        if len(args) == 1 and isinstance(args[0], (list, tuple)):
+            input_tensors = list(args[0])
+        else:
+            input_tensors = list(args)
+
+        dim = kwargs.get("dim", 0)
+
+    return input_tensors, dim
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.cat.default,
+    supports_dynamic_shapes=True,
+)
 def aten_ops_cat(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    inputs, dim = parse_cat_args(args, kwargs)
     return impl.cat.cat(
         ctx,
         target,
         SourceIR.ATEN,
         name,
-        input=args[0],
-        dim=args_bounds_check(args, 1, 0),
+        input=inputs,
+        dim=dim,
     )
 
 

tests/py/dynamo/conversion/test_cat_aten.py

@@ -25,6 +25,41 @@ def forward(self, x, y, z):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            ("pos", 1),
+            ("neg", -2),
+        ]
+    )
+    def test_cat_dim_in_kwargs(self, _, dim):
+        class Cat(nn.Module):
+            def forward(self, x, y, z):
+                return torch.ops.aten.cat.default((x, y, z), dim=dim)
+
+        inputs = [torch.randn(1, 2, 3), torch.randn(1, 1, 3), torch.randn(1, 3, 3)]
+        self.run_test(
+            Cat(),
+            inputs,
+        )
+
+    @parameterized.expand(
+        [
+            ("pos", 0),
+            ("neg", -3),
+        ]
+    )
+    def test_cat_with_scalar_inputs(self, _, dim):
+        # Ensure scalar tensor wrap works
+        class Cat(nn.Module):
+            def forward(self, x, y):
+                # y is a scalar, x is a tensor
+                return torch.ops.aten.cat.default((x, y), dim)
+
+        x = torch.randn(1, 2, 3, device="cuda")
+        y = torch.ones_like(x) * 5.0  # simulate scalar broadcast
+        inputs = [x, y]
+        self.run_test(Cat(), inputs)
+
     @parameterized.expand(
         [
             ("pos", 1),