Repository: pytorch/TensorRT

Issue:
🐛 [Bug] `aten.mean.dim` converter throws error despite being supported
##  Bug Description

The `aten.mean.dim` converter throws the following error when compiling the displayed model:

```python
##### MODEL:
class Sample(torch.nn.Module):
    def __init__(self):
        super(Sample, self).__init__()

    def forward(self, x):
        return torch.mean(x, dim=1)

##### ERROR:
  File "~/TensorRT/py/torch_tensorrt/fx/fx2trt.py", line 328, in call_function
    return converter(self.network, target, args, kwargs, self._cur_node_name)
  File "~/TensorRT/py/torch_tensorrt/fx/converters/aten_ops_converters.py", line 57, in aten_ops_adaptive_avg_poolnd
    raise RuntimeError(f"We do not support {target} has dim={args[1]}")
RuntimeError: We do not support aten.mean.dim has dim=[1]
```

## To Reproduce

Steps to reproduce the behavior:

1. Initialize model as above: `Sample().eval().cuda()`
2. Initialize oneinput tensors, for example: `torch.zeros((5, 5), dtype=torch.float, device="cuda:0")`
3. Compile the model using FX: `torch_tensorrt.fx.compile(model, [input_], min_acc_module_size=1, is_aten=True)`

## Expected behavior

Model should compile via the FX path or list the operator as unsupported.

## Environment

- Transformers: 4.26.1
 - Torch-TensorRT Version (e.g. 1.0.0): fce0a01e
 - PyTorch Version (e.g. 1.0): 2.1.0.dev20230313+cu117
- CPU Architecture: Intel Xeon CPU
- OS: Ubuntu 20.04
- How you installed PyTorch: pip
- Build command you used: `python setup.py develop`
- Are you using local sources or building from archives: local
- Python version: 3.8.13
- CUDA version: 11.7

## Additional Context
Solving this issue will also resolve the error encountered in #1740 


Discussion:
Addressed by #1657, but seeing the following error message when using that PR:

<details>

<summary>
<b>Error Message </b>

</summary>

```python
Got 1 acc subgraphs and 0 non-acc subgraphs
Traceback (most recent call last):
  File "case_test.py", line 22, in <module>
    main()
  File "case_test.py", line 19, in main
    trt_model = torch_tensorrt.fx.compile(model, [input_], min_acc_module_size=1, is_aten=True)
  File "~/TensorRT/py/torch_tensorrt/fx/lower.py", line 86, in compile
    return lowerer(module, input)
  File "~/TensorRT/py/torch_tensorrt/fx/lower.py", line 316, in __call__
    return do_lower(module, inputs)
  File "~/TensorRT/py/torch_tensorrt/fx/passes/pass_utils.py", line 118, in pass_with_validation
    processed_module = pass_(module, input, *args, **kwargs)
  File "~/TensorRT/py/torch_tensorrt/fx/lower.py", line 313, in do_lower
    lower_result = pm(module)
  File "/usr/local/lib/python3.8/dist-packages/torch/fx/passes/pass_manager.py", line 246, in __call__
    out = _pass(out)
  File "/usr/local/lib/python3.8/dist-packages/torch/fx/passes/pass_manager.py", line 246, in __call__
    out = _pass(out)
  File "~/TensorRT/py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py", line 202, in lower_func
    lowered_module = self._lower_func(
  File "~/TensorRT/py/torch_tensorrt/fx/lower.py", line 178, in lower_pass
    interp_res: TRTInterpreterResult = interpreter(mod, input, module_name)
  File "~/TensorRT/py/torch_tensorrt/fx/lower.py", line 130, in __call__
    interp_result: TRTInterpreterResult = interpreter.run(
  File "~/TensorRT/py/torch_tensorrt/fx/fx2trt.py", line 204, in run
    super().run()
  File "/usr/local/lib/python3.8/dist-packages/torch/fx/interpreter.py", line 137, in run
    self.env[node] = self.run_node(node)
  File "~/TensorRT/py/torch_tensorrt/fx/fx2trt.py", line 275, in run_node
    trt_node = super().run_node(n)
  File "/usr/local/lib/python3.8/dist-packages/torch/fx/interpreter.py", line 179, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "~/TensorRT/py/torch_tensorrt/fx/fx2trt.py", line 328, in call_function
    return converter(self.network, target, args, kwargs, self._cur_node_name)
  File "~/TensorRT/py/torch_tensorrt/fx/converters/fx2trt_ops_converter.py", line 81, in convert_avg_pool
    input_trt = add_missing_trt_tensors(network, [input])[0]
  File "~/TensorRT/py/torch_tensorrt/fx/converters/fx2trt_ops_converter_utils.py", line 19, in add_missing_trt_tensors
    dtype = check_torch_dtype(*tensors)
  File "~/TensorRT/py/torch_tensorrt/fx/converters/fx2trt_ops_converter_utils.py", line 88, in check_torch_dtype
    assert (
AssertionError: While executing %mean_dim : [#users=1] = call_function[target=torch.ops.aten.mean.dim](args = (%arg0, [1]), kwargs = {_itensor_to_tensor_meta: {<tensorrt.tensorrt.ITensor object at 0x7f48305059f0>: None}})
Original traceback:
  File "case_test.py", line 13, in forward
    return torch.mean(x, dim=1)
```

</details>

Resolving patch:
diff --git a/py/torch_tensorrt/fx/converters/aten_ops_converters.py b/py/torch_tensorrt/fx/converters/aten_ops_converters.py
--- a/py/torch_tensorrt/fx/converters/aten_ops_converters.py
+++ b/py/torch_tensorrt/fx/converters/aten_ops_converters.py
@@ -41,7 +41,6 @@ def aten_ops_add(
     return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.mean.dim)
 @tensorrt_converter(torch.ops.aten._adaptive_avg_pool3d.default)
 @tensorrt_converter(torch.ops.aten._adaptive_avg_pool2d.default)
 def aten_ops_adaptive_avg_poolnd(
@@ -51,24 +50,38 @@ def aten_ops_adaptive_avg_poolnd(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    if target == torch.ops.aten.mean.dim:
-
-        if list(args[1]) != [-1, -2]:
-            raise RuntimeError(f"We do not support {target} has dim={args[1]}")
-        else:
-            output_size = [1, 1]
-    else:
-        output_size = args[1]
-
     kwargs_new = {
         "input": args[0],
-        "output_size": output_size,
+        "output_size": args[1],
     }
     return acc_ops_converters.acc_ops_adaptive_avg_poolnd(
         network, target, None, kwargs_new, name
     )
 
 
+@tensorrt_converter(torch.ops.aten.mean.default)
+@tensorrt_converter(torch.ops.aten.mean.dim)
+def aten_ops_mean(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> TRTTensor:
+    # Default invocation of aten.mean only uses first argument and
+    # averages over all elements (all dimensions)
+    # aten.mean.dim invocation allows specification of dimensions to average
+    # over, as well at the option to keep the dimension or not
+    kwargs_new = {
+        "input": args[0],
+        "dim": args[1] if len(args) >= 2 else list(range(len(args[0].shape))),
+        "keepdim": args[2] if len(args) >= 3 else False,
+    }
+    return add_reduce_layer(
+        network, target, args, kwargs_new, trt.ReduceOperation.AVG, name
+    )
+
+
 @tensorrt_converter(torch.ops.aten.batch_norm)
 def aten_ops_batch_norm(
     network: TRTNetwork,


Test patch:
diff --git a/py/torch_tensorrt/fx/test/converters/aten_op/test_mean_aten.py b/py/torch_tensorrt/fx/test/converters/aten_op/test_mean_aten.py
new file mode 100644
--- /dev/null
+++ b/py/torch_tensorrt/fx/test/converters/aten_op/test_mean_aten.py
@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestMeanDimConverter(DispatchTestCase):
+    def test_mean_dim_keepdims(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x, dim=[0, 1], keepdim=True)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.mean.dim})
+
+    def test_mean_dim_keepdims_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x, dim=[0, 1, 2], keepdim=True)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.mean.dim}
+        )
+
+    def test_mean_dim_keepdims_false(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x, dim=0, keepdim=False)
+
+        inputs = [torch.randn(3, 5, 7)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.mean.dim})
+
+    def test_mean_dim_keepdims_false_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x, dim=-1, keepdim=False)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.mean.dim}
+        )
+
+
+class TestMeanConverter(DispatchTestCase):
+    def test_mean(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x)
+
+        inputs = [torch.randn(3, 8, 5, 7, 1)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.mean.default})
+
+    def test_mean_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return torch.mean(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 5, 8), (3, 10, 10))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.mean.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()