[torchlib] Add missing dtype parameter to aten_mean_dim (#2885)

Fixes #2884

`aten_mean_dim` and `aten_mean_dim_complex` are missing the `dtype`
keyword argument from their signatures, even though the ATen schema
documents it (`ScalarType? dtype=None`). This causes a `TypeError` when
PyTorch lowers `aten::mean.dim` with an explicit `dtype` - which happens
for any model using `GlobalAveragePooling2D` (Keras/PyTorch).

- Add `dtype: int = -1` to `aten_mean_dim`, with `op.Cast` when dtype is
specified
- Add `dtype: int = -1` to `aten_mean_dim_complex`, raising
`NotImplementedError` for complex tensors

Follows the same pattern used by `aten_sum_dim_IntList` and
`aten_sum_dim_IntList_complex`.

Author: linusjuni

onnxscript/function_libs/torch_lib/ops/core.py

@@ -6231,19 +6231,25 @@ def aten_mean_complex(self: TReal) -> TReal:
 
 
 @torch_op("aten::mean.dim", trace_only=True)
-def aten_mean_dim(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
+def aten_mean_dim(self: TReal, dim: INT64, keepdim: bool = False, dtype: int = -1) -> TReal:
     """mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor"""
 
     if len(self.shape) == 0:
         result = self
     else:
         dims = op.Reshape(dim, op.Constant(value_ints=[-1]))
         result = op.ReduceMean(self, dims, keepdims=keepdim)
+
+    if dtype != -1 and dtype is not None:
+        result = op.Cast(result, to=dtype)
+
     return result
 
 
 @torch_op("aten::mean.dim", trace_only=True, complex=True)
-def aten_mean_dim_complex(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
+def aten_mean_dim_complex(
+    self: TReal, dim: INT64, keepdim: bool = False, dtype: int = -1
+) -> TReal:
     """mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor"""
 
     if len(self.shape) == 1:
@@ -6254,6 +6260,12 @@ def aten_mean_dim_complex(self: TReal, dim: INT64, keepdim: bool = False) -> TRe
         dim = op.Where(op.Less(dim, zero), op.Sub(dim, one), dim)
         dims = op.Reshape(dim, op.Constant(value_ints=[-1]))
         result = op.ReduceMean(self, dims, keepdims=keepdim)
+
+    if dtype != -1 and dtype is not None:
+        raise NotImplementedError(
+            "support for the dtype argument is not implemented for complex tensors"
+        )
+
     return result