Remove the MLIR passes that are Migrated to tflite.

PiperOrigin-RevId: 663611553
google-ai-edge · Aug 16, 2024 · a54d10d · a54d10d
1 parent 7e6023e
commit a54d10d
Show file tree

Hide file tree

Showing 8 changed files with 288 additions and 134 deletions.
diff --git a/ai_edge_torch/_convert/fx_passes/test/test_build_aten_composite_pass.py b/ai_edge_torch/_convert/fx_passes/test/test_build_aten_composite_pass.py
@@ -65,16 +65,24 @@ def test_hardswish_layer(self):
     stablehlo = _export_to_stablehlo_with_composite(
         lambda x: torch.nn.Hardswish()(x), (torch.rand(10, 10),)  # pylint: disable=unnecessary-lambda
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.hardswish.default"'), 1
+
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.hardswish.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_hardswish_op(self):
     stablehlo = _export_to_stablehlo_with_composite(
         lambda x: torch.ops.aten.hardswish.default(x), (torch.rand(10, 10),)  # pylint: disable=unnecessary-lambda
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.hardswish.default"'), 1
+
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.hardswish.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_avg_pool2d_layer(self):
@@ -90,8 +98,11 @@ def test_avg_pool2d_layer(self):
         )(x),
         (torch.rand(1, 3, 6, 6),),
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.avg_pool2d.default"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.avg_pool2d.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_avg_pool2d_op(self):
@@ -108,8 +119,11 @@ def test_avg_pool2d_op(self):
         ),
         (torch.rand(1, 3, 6, 6),),
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.avg_pool2d.default"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.avg_pool2d.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_avg_pool2d_ceil_mode(self):
@@ -126,41 +140,56 @@ def test_avg_pool2d_ceil_mode(self):
         ),
         (torch.rand(1, 3, 6, 6),),
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.avg_pool2d.default"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.avg_pool2d.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_gelu_layer(self):
     stablehlo = _export_to_stablehlo_with_composite(
         lambda x: torch.nn.GELU()(x), (torch.rand(10, 10),)  # pylint: disable=unnecessary-lambda
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.gelu.default"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.gelu.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_approximate_gelu_layer(self):
     stablehlo = _export_to_stablehlo_with_composite(
         lambda x: torch.nn.GELU('tanh')(x), (torch.rand(10, 10),)  # pylint: disable=unnecessary-lambda
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "aten.gelu.default"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "aten.gelu.default"': 1},
+        {'stablehlo.custom_call @mark_tensor': 2},
     )
 
   def test_embedding_lookup_layer(self):
     stablehlo = _export_to_stablehlo_with_composite(
         torch.nn.Embedding(10, 10), (torch.full((1, 10), 0, dtype=torch.long),)
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "odml.embedding_lookup"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "odml.embedding_lookup"': 1},
+        {'stablehlo.custom_call @mark_tensor': 3},
     )
 
   def test_embedding_lookup_op(self):
     stablehlo = _export_to_stablehlo_with_composite(
         lambda *x: torch.ops.aten.embedding.default(*x),
         (torch.rand(10, 10), torch.full((1, 10), 0, dtype=torch.long)),
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "odml.embedding_lookup"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "odml.embedding_lookup"': 1},
+        {'stablehlo.custom_call @mark_tensor': 3},
     )
 
   def test_embedding_lookup_functional(self):
@@ -171,8 +200,11 @@ def test_embedding_lookup_functional(self):
             torch.rand(10, 10),
         ),
     )
-    self.assertEqual(
-        stablehlo.count('stablehlo.composite "odml.embedding_lookup"'), 1
+    lowertools.assert_string_count(
+        self,
+        stablehlo,
+        {'stablehlo.composite "odml.embedding_lookup"': 1},
+        {'stablehlo.custom_call @mark_tensor': 3},
     )