fix: Allow rank differences in aten.expand

- Add support for `aten.expand.default` in Dynamo converter registry
- Build converter to support rank-padding for input Tensors, in line
with the existing Torch behavior
- Add test case to validate new behavior, in addition to existing cases
validating old behavior

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
+import tensorrt as trt
 import torch
 from torch.fx.node import Argument, Node, Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
@@ -12,8 +13,6 @@
 from torch_tensorrt.fx.converters import acc_ops_converters
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-import tensorrt as trt
-
 from .converter_registry import dynamo_tensorrt_converter
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -380,3 +379,21 @@ def aten_ops_permute(
         args[0],
         args[1],
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.expand.default)
+def aten_ops_expand(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unsqueeze.expand(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )

py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py

@@ -5,6 +5,7 @@
 from torch_tensorrt.fx.converters.converter_utils import (
     get_positive_dim,
     get_trt_tensor,
+    prepend_ones,
     set_layer_name,
 )
 from torch_tensorrt.fx.types import Shape, TRTNetwork, TRTTensor
@@ -27,7 +28,7 @@ def unsqueeze(
         )
 
     dim = cast(int, dim)
-    input_shape = input_val.shape
+
     input_shape_size = (
         len(input_val.shape) + 1
         if network.has_implicit_batch_dimension
@@ -46,5 +47,55 @@ def unsqueeze(
     layer.reshape_dims = (
         tuple(input_val.shape)[:dim] + (1,) + tuple(input_val.shape)[dim:]
     )
-    set_layer_name(layer, target, name)
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)
+
+
+def expand(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_t: TRTTensor,
+    shape: Shape,
+) -> TRTTensor:
+    if not isinstance(input_t, TRTTensor):
+        raise RuntimeError(
+            f"expand received input {input_t} that is not a TensorRT ITensor"
+        )
+
+    shape_rank = len(shape)
+    initial_tensor_rank = len(input_t.shape)
+
+    # If the rank of the input tensor is less than the shape's rank, pad with ones
+    if initial_tensor_rank < shape_rank:
+        input_t = prepend_ones(
+            network,
+            input_t,
+            name + "_expand_ones_padding",
+            shape_rank - initial_tensor_rank,
+        )
+    # If the rank of the input tensor is more than the shape's rank, raise error
+    elif initial_tensor_rank > shape_rank:
+        raise RuntimeError(
+            f"expand called with {shape_rank}-dimensional shape on Tensor with {len(shape)} dimensions. "
+            "Cannot expand to shape with rank smaller than original tensor."
+        )
+
+    # After the above padding, the shape and tensor rank must be equal
+    assert len(input_t.shape) == shape_rank
+
+    # -1 denotes taking the shape from the original input tensor
+    shape = tuple(
+        [input_t.shape[i] if shape[i] == -1 else shape[i] for i in range(shape_rank)]
+    )
+
+    # Establish the desired output shape, strides, and starting indices
+    input_tensor_shape = tuple(input_t.shape)
+    start = tuple([0] * shape_rank)
+    stride = tuple(
+        [int(i == o) for i, o in zip(input_tensor_shape, shape)]
+    )  # stride == 1 if dimensions match, 0 otherwise
+    layer = network.add_slice(input_t, start=start, shape=shape, stride=stride)
+    set_layer_name(layer, target, name, source_ir)
     return layer.get_output(0)

tests/py/dynamo/converters/test_expand_aten.py

@@ -1,8 +1,8 @@
 import torch
 import torch.nn as nn
+from harness import DispatchTestCase
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
-from harness import DispatchTestCase
 
 
 class TestExpandConverter(DispatchTestCase):
@@ -12,6 +12,7 @@ class TestExpandConverter(DispatchTestCase):
             ("3d_dim", (2, 3, 4), (2, 1, 1)),
             ("4d_dim", (2, 3, 4, 5), (2, 1, 1, 1)),
             ("keep_dim", (2, 3, -1, -1), (2, 1, 5, 5)),
+            ("different_ranks", (2, 3, -1, -1), (1, 5, 7)),
         ]
     )
     def test_expand(self, _, sizes, init_size):