Merge branch 'main' into dev/gokul/fix_div

Author: gokulkrishna98

tests/test_stateful.py

@@ -26,7 +26,6 @@
 from .utils import (
     filecheck_pattern,
     make_dynamic_shapes,
-    run_transforms,
     validate_numerical_output,
 )
 
@@ -229,7 +228,7 @@ async def test_buffer_mutation_after_make_stateful(self) -> None:
             .add_exported_program(self._mutable_buffer_program())
             .to_coreai()
         )
-        await run_transforms(result)
+        result.optimize()
         filecheck_pattern(
             str(result),
             check_file="""
@@ -395,7 +394,7 @@ async def test_user_input_mutation_after_make_stateful(self) -> None:
         result = (
             TorchConverter().add_exported_program(self._mutating_program()).to_coreai()
         )
-        await run_transforms(result)
+        result.optimize()
         filecheck_pattern(
             str(result),
             check_file="""
@@ -507,7 +506,7 @@ async def test_both_mutations_after_make_stateful(self) -> None:
             .add_exported_program(self._both_mutations_program())
             .to_coreai()
         )
-        await run_transforms(result)
+        result.optimize()
         filecheck_pattern(
             str(result),
             check_file="""

tests/utils.py

@@ -16,9 +16,6 @@
 import numpy as np
 import numpy.typing as npt
 import torch
-from coreai._compiler._transforms import GlobalOptions, PassEntry, apply_passes
-from coreai._compiler._transforms.passes import CorePasses
-from coreai.authoring import AIProgram
 from coreai.runtime import NDArray, StorageKind
 from filecheck.matcher import Matcher
 from filecheck.options import Options
@@ -81,19 +78,6 @@ def _get_test_specialization_options() -> "SpecializationOptions | None":
     raise ValueError(msg)
 
 
-async def run_transforms(coreai_program: AIProgram) -> None:
-    """Run essential transformation passes."""
-    await apply_passes(
-        coreai_program._mlir_module,
-        passes=[
-            PassEntry.get(CorePasses._CORE_OPTIMIZE),
-            PassEntry.get(CorePasses._UPDATE_SIGNATURE_TO_HANDLES),
-            PassEntry.get(CorePasses._PROPAGATE_HANDLE_UPDATES),
-        ],
-        options=GlobalOptions(Path()),
-    )
-
-
 def make_dynamic_shapes(
     **arg_specs: list[str | None] | dict[int, str | None],
 ) -> dict[str, dict[int, torch.export.Dim]]: