Add impl for aten.mean.dim and lowering.

PiperOrigin-RevId: 768227681

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -110,6 +110,11 @@ def _aten_bitwise_and_tensor_decomp(x, y):
   return torch.ops.tfl.logical_and(x, y)
 
 
+@register_decomp(torch.ops.aten.mean.dim)
+def _aten_mean_dim_decomp(x, dim, keepdim=False):
+  return torch.ops.tfl.mean(x, dim, keepdim)
+
+
 @register_decomp(torch.ops.aten.gt.Tensor)
 def _aten_gt_tensor_decomp(x, y):
   return torch.ops.tfl.greater(x, y)

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -177,6 +177,27 @@ def _tfl_logical_and_lowering(
   )
 
 
+@lower(torch.ops.tfl.mean.default)
+def _tfl_mean_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    dims: int | ir.Value | Sequence[int | ir.Value],
+    keepdim: bool = False,
+) -> ir.Value:
+  if isinstance(dims, int) or isinstance(dims, ir.Value):
+    dims_ir_value = lowering_utils.convert_to_ir_value(dims)
+  else:
+    dims_ir_value = lowering_utils.convert_shape_to_ir_value(dims)
+  return _ir_operation(
+      "tfl.mean",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x, dims_ir_value],
+      attributes={
+          "keep_dims": ir.BoolAttr.get(keepdim),
+      },
+  )
+
+
 @lower(torch.ops.tfl.greater.default)
 def _tfl_greater_lowering(
     lctx: LoweringContext,

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -68,6 +68,13 @@ def tfl_logical_and(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
   return torch.logical_and(x, y)
 
 
+@custom_op_with_fake(
+    "tfl::mean", schema="(Tensor x, Any dims, bool keepdim) -> Tensor"
+)
+def tfl_mean(x: torch.Tensor, dims: Any, keepdim: bool = False) -> torch.Tensor:
+  return torch.mean(x, dims, keepdim)
+
+
 @custom_op_with_fake("tfl::greater")
 def tfl_greater(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
   return torch.gt(x, y)

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -134,6 +134,9 @@ def _assert_export_and_close(
       ("aten_pow_Tensor_Scalar_0", torch.ops.aten.pow.Tensor_Scalar, (rnd(torch.float32, (10, 10)), np.random.rand(),), dict()),
       ("aten_pow_Tensor_Tensor_0", torch.ops.aten.pow.Tensor_Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
       ("aten_bitwise_and_Tensor_0", torch.ops.aten.bitwise_and.Tensor, (rnd(torch.bool, (10, 10)), rnd(torch.bool, (10, 10)),), dict()),
+      ("aten_mean_dim_0", torch.ops.aten.mean.dim, (rnd(torch.float32, (10, 10)), 0), dict()),
+      ("aten_mean_dim_1", torch.ops.aten.mean.dim, (rnd(torch.float32, (10, 10)), 0, True), dict()),
+      ("aten_mean_dim_2", torch.ops.aten.mean.dim, (rnd(torch.float32, (10, 10)), 1), dict()),
       ("aten_gt_Tensor_0", torch.ops.aten.gt.Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
       ("aten_gt_Tensor_1", torch.ops.aten.gt.Tensor, (rnd(torch.float32, (1, 10)), rnd(torch.float32, (10, 1)),), dict()),
       ("aten_lt_Tensor_0", torch.ops.aten.lt.Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),