aten::cat converter moving to impl

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -46,6 +46,24 @@ def aten_ops_batch_norm(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.cat.default)
+def aten_ops_cat(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.cat.cat(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        tensors=args[0],
+        dim = args_bounds_check(args, 2, 1),
+    )
+
+
 def embedding_param_validator(embedding_node: Node) -> bool:
     scale_grad_by_freq = args_bounds_check(embedding_node.args, 3)
     sparse = args_bounds_check(embedding_node.args, 4)

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -3,6 +3,7 @@
 from . import (
     activation,
     cast,
+    cat,
     condition,
     conv,
     elementwise,

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

@@ -0,0 +1,29 @@
+from typing import Optional, Union, Sequence, Dict
+
+import torch
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+def cat(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTNetwork,
+    dim: int,
+    
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+
+    if any(not isinstance(t, TRTTensor) for t in input):  # type: ignore[union-attr]
+        raise RuntimeError(
+            f"cat received inputs {input} that is not part " "of the TensorRT region!"
+        )
+    concat_layer = network.add_concatenation(input)
+    if dim < 0:
+        dim = len(input[0].shape) + dim
+
+    concat_layer.axis = dim
+    set_layer_name(concat_layer, target, name + "_gather", source_ir)
+    return concat_layer.get_output(0)