Skip to content

Refactor test models for ort_fusions #2237

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
merged 3 commits into from
Apr 28, 2025
Merged

Refactor test models for ort_fusions #2237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0fb399f
refactor model store
shubhambhokare1 Apr 28, 2025
77f0e00
remove extra import
shubhambhokare1 Apr 28, 2025
6264674
lint fixes
shubhambhokare1 Apr 28, 2025
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
8 changes: 4 additions & 4 deletions onnxscript/rewriter/ort_fusions/cos_sin_cache_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,14 @@
from parameterized import parameterized

import onnxscript.optimizer
from onnxscript.rewriter.ort_fusions._rotary_embedding_models import (
from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
from onnxscript.rewriter.ort_fusions.models._rotary_embedding_models import (
partial_rotary_test_case,
test_case_1,
test_case_2,
)
from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
from onnxscript.rewriter.ort_fusions.models._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions.rotary_embedding import (
fuse_partial_rotary_embedding,
fuse_rotary_embedding,
Expand Down
2 changes: 1 addition & 1 deletion onnxscript/rewriter/ort_fusions/fuse_xformers_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@

import onnxscript.optimizer
from onnxscript.rewriter.ort_fusions._core import fuse_xformers
from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.models._smollm_1 import smollm_test_1


class TestFuseXformers(unittest.TestCase):
Expand Down
2 changes: 1 addition & 1 deletion onnxscript/rewriter/ort_fusions/mha_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@

import onnxscript.optimizer
import onnxscript.rewriter.ort_fusions._core as xformers
from onnxscript.rewriter.ort_fusions._smollm_2 import smollm_test_2
from onnxscript.rewriter.ort_fusions._test_utils import ORT_VERSION, assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.models._smollm_2 import smollm_test_2


class TestMultiHeadAttention(unittest.TestCase):
Expand Down
0 ...r/ort_fusions/_rotary_embedding_models.py → ...usions/models/_rotary_embedding_models.py
View file
Open in desktop
File renamed without changes.
0 onnxscript/rewriter/ort_fusions/_smollm_1.py → .../rewriter/ort_fusions/models/_smollm_1.py
View file
Open in desktop
File renamed without changes.
0 onnxscript/rewriter/ort_fusions/_smollm_2.py → .../rewriter/ort_fusions/models/_smollm_2.py
View file
Open in desktop
File renamed without changes.
0 ...ript/rewriter/ort_fusions/_test_models.py → ...writer/ort_fusions/models/_test_models.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion onnxscript/rewriter/ort_fusions/rms_normalization_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import unittest

import onnxscript.optimizer
from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.models._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions.rms_normalization import fuse_rms_normalization


Expand Down
4 changes: 2 additions & 2 deletions onnxscript/rewriter/ort_fusions/rotary_embedding_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
from parameterized import parameterized

import onnxscript.optimizer
from onnxscript.rewriter.ort_fusions._rotary_embedding_models import test_case_1
from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions.models._rotary_embedding_models import test_case_1
from onnxscript.rewriter.ort_fusions.models._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions.rotary_embedding import fuse_rotary_embedding


Expand Down
2 changes: 1 addition & 1 deletion onnxscript/rewriter/ort_fusions/skip_normalization_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import unittest

import onnxscript.optimizer
from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
from onnxscript.rewriter.ort_fusions.models._smollm_1 import smollm_test_1
from onnxscript.rewriter.ort_fusions.rms_normalization import fuse_rms_normalization
from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_skip_rms_normalization

Expand Down
Loading