Add impl for aten.mean.dim and lowering.

PiperOrigin-RevId: 764338440

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)
@@ -203,6 +208,19 @@ def _aten_cat_decomp(tensors, dim=0):
   return torch.ops.tfl.concatenation(processed_tensors, dim)
 
 
+@register_decomp(torch.ops.aten.full_like.default)
+def _aten_full_like_decomp(
+    x,
+    fill_value,
+    dtype=None,
+    layout=None,
+    device=None,
+    pin_memory=None,
+    memory_format=None,
+):
+  return torch.ops.tfl.fill(tuple(x.shape), fill_value)
+
+
 @register_decomp(torch.ops.aten.view.default)
 def _aten_view_decomp(x, shape: Sequence[int]):
   return torch.ops.tfl.reshape(x, shape)

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,
@@ -317,45 +338,69 @@ def _tfl_concatenation_lowering(
   )
 
 
+@lower(torch.ops.tfl.fill.default)
+def _tfl_fill_lowering(
+    lctx: LoweringContext,
+    dims: Sequence[int | ir.Value],
+    fill_value: ir.Value,
+) -> ir.Value:
+  dims_ir_value = lowering_utils.convert_shape_to_ir_value(dims)
+  fill_value_ir_value = lowering_utils.convert_to_ir_value(fill_value)
+
+  # Ensure fill_value_ir_value is a scalar (0-D tensor) for TFLite Fill op.
+  # The TFLite Fill kernel expects the value to be a 0-D tensor.
+  if isinstance(fill_value_ir_value.type, ir.RankedTensorType):
+    tensor_type = fill_value_ir_value.type
+    # If it's a 1-D tensor with a single element, reshape to 0-D.
+    if list(tensor_type.shape) == [1]:
+      scalar_type = ir.RankedTensorType.get([], tensor_type.element_type)
+      fill_value_ir_value = stablehlo.reshape(scalar_type, fill_value_ir_value)
+
+  # Determine the target element type from the node's output definition.
+  result_types = lowering_utils.node_meta_to_ir_types(lctx.node)
+  if not result_types or not isinstance(result_types[0], ir.RankedTensorType):
+    raise ValueError(
+        "tfl.fill: Unable to determine result tensor type or result is not a"
+        " ranked tensor."
+    )
+  target_element_type = result_types[0].element_type
+
+  # Ensure fill_value_ir_value is a RankedTensorType to access its properties.
+  if not isinstance(fill_value_ir_value.type, ir.RankedTensorType):
+    raise TypeError(
+        "tfl.fill: fill_value_ir_value expected to be RankedTensorType, got"
+        f" {fill_value_ir_value.type}"
+    )
+
+  current_fill_tensor_type = fill_value_ir_value.type
+  current_element_type = current_fill_tensor_type.element_type
+
+  # If the element type of the (scalar) fill_value doesn't match the target
+  # output element type, cast fill_value_ir_value to the target_element_type
+  # while maintaining its current shape (which should be scalar).
+  if current_element_type != target_element_type:
+    cast_to_type = ir.RankedTensorType.get(
+        current_fill_tensor_type.shape, target_element_type
+    )
+    fill_value_ir_value = stablehlo.convert(cast_to_type, fill_value_ir_value)
+
+  return _ir_operation(
+      "tfl.fill",
+      results=result_types,
+      operands=[dims_ir_value, fill_value_ir_value],
+  )
+
+
 @lower(torch.ops.tfl.reshape.default)
 def _tfl_reshape_lowering(
     lctx: LoweringContext,
     x: ir.Value,
     shape: Sequence[int | ir.Value],
 ) -> ir.Value:
-  # Check if all elements in the shape sequence are integers.
-  if not shape or all(isinstance(dim, int) for dim in shape):
-    # If all are integers, create a constant numpy array.
-    # Assuming int32 is the required type for TFLite shape tensors.
-    shape_ir_value = lowering_utils.numpy_array_constant(
-        np.array(shape, dtype=np.int32)
-    )
-  else:
-    # Handle mixed int and ir.Value shape sequence
-    processed_dims = []
-    for dim in shape:
-      if isinstance(dim, int):
-        # Convert int to a constant 1D tensor
-        shape_ir_value = lowering_utils.numpy_array_constant(
-            np.array([dim], dtype=np.int32)
-        )
-        processed_dims.append(shape_ir_value)
-      else:
-        assert isinstance(dim, ir.Value)
-        # Convert ir.Value to a constant 1D tensor
-        new_type = ir.RankedTensorType.get([1], dim.type.element_type)
-        reshape_dim = stablehlo.reshape(new_type, dim)
-        processed_dims.append(reshape_dim)
-
-    shape_ir_value = stablehlo.concatenate(
-        processed_dims,
-        dimension=0,
-    )
-
   return _ir_operation(
       "tfl.reshape",
       results=lowering_utils.node_meta_to_ir_types(lctx.node),
-      operands=[x, shape_ir_value],
+      operands=[x, lowering_utils.convert_shape_to_ir_value(shape)],
   )
 
 

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)
@@ -123,6 +130,11 @@ def tfl_concatenation(
   return torch.cat(tensors, dim=dim)
 
 
+@custom_op_with_fake("tfl::fill", schema="(int[] x, Any y) -> Tensor")
+def tfl_fill(dims: Sequence[int], fill_value: Any) -> torch.Tensor:
+  return torch.full(dims, fill_value)
+
+
 def _normalize_shape(
     tensor_input: torch.Tensor, shape: Sequence[int]
 ) -> Sequence[int]:

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()),
@@ -163,6 +166,8 @@ def _assert_export_and_close(
       ("aten_cat_2", torch.ops.aten.cat.default, ([rnd(torch.float32, (10, 10)), rnd(torch.float32, (0, 10))], 0,), dict()),
       ("aten_cat_3", torch.ops.aten.cat.default, ([rnd(torch.float32, (10, 10)), rnd(torch.float32, (0,))], 0,), dict()),
       ("aten_cat_4", torch.ops.aten.cat.default, ([rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10))],), dict()),
+      ("aten_full_like_0", torch.ops.aten.full_like.default, (rnd(torch.float32, (10, 10)), 0.123,), dict()),
+      ("aten_full_like_1", torch.ops.aten.full_like.default, (rnd(torch.int64, (10, 10)), 123,), dict()),
       ("aten_view_0", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [1, 100],), dict()),
       ("aten_view_1", torch.ops.aten.view.default, (rnd(torch.float32, (1, 10)), [10, 1],), dict()),
       ("aten_view_2", torch.ops.aten.view.default, (rnd(torch.float32, (10, 10)), [2, 5, 10],), dict()),

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -17,7 +17,7 @@
 from collections.abc import Callable
 import functools
 import numbers
-from typing import Any, Optional, Union
+from typing import Any, Optional, Sequence, Union
 from ai_edge_torch.odml_torch import export_utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
@@ -281,3 +281,33 @@ def convert_to_ir_value(
   if isinstance(value, ir.Value):
     return value
   raise TypeError(f"Unsupported type for conversion to ir.Value: {type(value)}")
+
+
+def convert_shape_to_ir_value(
+    shape: Sequence[int],
+) -> ir.Value:
+  # Check if all elements in the shape sequence are integers.
+  if not shape or all(isinstance(dim, int) for dim in shape):
+    # If all are integers, create a constant numpy array.
+    # Assuming int32 is the required type for TFLite shape tensors.
+    shape_ir_value = numpy_array_constant(np.array(shape, dtype=np.int32))
+  else:
+    # Handle mixed int and ir.Value shape sequence
+    processed_dims = []
+    for dim in shape:
+      if isinstance(dim, int):
+        # Convert int to a constant 1D tensor
+        shape_ir_value = numpy_array_constant(np.array([dim], dtype=np.int32))
+        processed_dims.append(shape_ir_value)
+      else:
+        assert isinstance(dim, ir.Value)
+        # Convert ir.Value to a constant 1D tensor
+        new_type = ir.RankedTensorType.get([1], dim.type.element_type)
+        reshape_dim = stablehlo.reshape(new_type, dim)
+        processed_dims.append(reshape_dim)
+
+    shape_ir_value = stablehlo.concatenate(
+        processed_dims,
+        dimension=0,
+    )
+  return shape_ir_value