Skip to content

[clean up] Remove run_transforms helper in favor of AIProgram.optimize() #50

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gokulkrishna98 merged 1 commit into from
Jul 23, 2026
Merged

[clean up] Remove run_transforms helper in favor of AIProgram.optimize() #50

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions tests/test_stateful.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
from .utils import (
filecheck_pattern,
make_dynamic_shapes,
run_transforms,
validate_numerical_output,
)

Expand Down Expand Up @@ -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="""
Expand Down Expand Up @@ -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="""
Expand Down Expand Up @@ -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="""
Expand Down
16 changes: 0 additions & 16 deletions tests/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]]:
Expand Down