Add impl for aten.squeeze.dims and lowering.

PiperOrigin-RevId: 773096398

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -248,6 +248,16 @@ def _aten_expand_decomp(x, shape: Sequence[int]):
   return torch.ops.tfl.broadcast_to(x, shape)
 
 
+@register_decomp(torch.ops.aten.squeeze.dims)
+def _aten_squeeze_dims_decomp(x, squeeze_dims: Sequence[int]):
+  if len(squeeze_dims) > 8:
+    raise ValueError(
+        "torch.ops.tfl.squeeze supports squeezing at most 8 dimensions, but got"
+        f" {len(squeeze_dims)} dimensions."
+    )
+  return torch.ops.tfl.squeeze(x, squeeze_dims)
+
+
 @register_decomp(torch.ops.aten._softmax.default)
 def _aten__softmax_decomp(
     x, dim: int, half_to_float: bool  # pylint: disable=unused-argument

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -525,6 +525,25 @@ def _tfl_broadcast_to_lowering(
   )
 
 
+@lower(torch.ops.tfl.squeeze.default)
+def _tfl_squeeze_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    squeeze_dims: Sequence[int | ir.Value],
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.squeeze",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x],
+      attributes={
+          "squeeze_dims": ir.ArrayAttr.get([
+              ir.IntegerAttr.get(ir.IntegerType.get_signless(64), int(d))
+              for d in squeeze_dims
+          ]),
+      },
+  )
+
+
 @lower(torch.ops.tfl.softmax.default)
 def _tfl_softmax_lowering(
     lctx: LoweringContext,

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -237,6 +237,11 @@ def tfl_broadcast_to(x: torch.Tensor, shape: Sequence[int]) -> torch.Tensor:
   return x.expand(shape).clone()
 
 
+@custom_op_with_fake("tfl::squeeze")
+def tfl_squeeze(x: torch.Tensor, squeeze_dims: Sequence[int]) -> torch.Tensor:
+  return torch.squeeze(x, squeeze_dims).clone()
+
+
 @custom_op_with_fake("tfl::softmax")
 def tfl_softmax(x: torch.Tensor) -> torch.Tensor:
   return torch.nn.functional.softmax(x, dim=-1)

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

@@ -182,6 +182,7 @@ def _assert_export_and_close(
       ("aten_unsqueeze_3", torch.ops.aten.unsqueeze.default, (rnd(torch.float32, (10, 10)), -1,), dict()),
       ("aten_expand_0", torch.ops.aten.expand.default, (rnd(torch.float32, (10, 1)), [10, 10],), dict()),
       ("aten_expand_1", torch.ops.aten.expand.default, (rnd(torch.float32, (1, 10)), [10, 10],), dict()),
+      ("aten_squeeze_dims_0", torch.ops.aten.squeeze.dims, (rnd(torch.float32, (2, 1, 2, 1, 2)), [1, 2, 3],), dict()),
       ("aten__softmax_0", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), -1, False), dict()),
       ("aten__softmax_1", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), -1, False), dict()),
       ("aten__softmax_2", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), 0, False), dict()),