Skip to content

Update test_export_fx_passes.py #8933

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

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Update test_export_fx_passes.py #8933

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
15dcf52
Update test_export_fx_passes.py
avikchaudhuri Apr 3, 2025
f815734
Create .torch_pin
avikchaudhuri Jun 18, 2025
dbb9af9
Update test_unbounded_dynamism.py
avikchaudhuri Jun 18, 2025
f81505c
Merge branch 'master' into pr/8933
avikchaudhuri Jun 18, 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
1 change: 1 addition & 0 deletions .torch_pin
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
#150583
4 changes: 2 additions & 2 deletions test/stablehlo/test_export_fx_passes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class ExportFxPassTest(unittest.TestCase):

def test_decompose_dynamic_shape_select(self):
args = (torch.rand((10, 197, 768)), 1, 0)
dynamic_shapes = (({0: Dim("bs")}, None, None),)
dynamic_shapes = ({0: Dim("bs")}, None, None)
m = wrap_func_as_nn_module(torch.ops.aten.select.int)
ep = export(m, args, dynamic_shapes=dynamic_shapes)
out1 = ep.module()(*args)
Expand Down Expand Up @@ -55,7 +55,7 @@ def forward(self, x):
def test_embedding_indices_flatten(self):
args = (torch.rand((20, 768)), torch.randint(0, 15,
(3, 10)).to(torch.int64))
dynamic_shapes = ((None, {0: Dim("bs")}),)
dynamic_shapes = (None, {0: Dim("bs")})
m = wrap_func_as_nn_module(torch.ops.aten.embedding.default)
ep = export(m, args, dynamic_shapes=dynamic_shapes)
print(ep)
Expand Down
52 changes: 26 additions & 26 deletions test/stablehlo/test_unbounded_dynamism.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ class UnboundedDynamismExportTest(unittest.TestCase):

def test_add(self):
args = (torch.rand((10, 197, 768)), torch.rand((10, 197, 768)))
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.add.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -45,7 +45,7 @@ def test_add(self):

def test_add_scalar(self):
args = (torch.rand((10, 197, 768)), 0.345)
dynamic_shapes = (({0: Dim("dim")}, None),)
dynamic_shapes = ({0: Dim("dim")}, None)
m = wrap_func_as_nn_module(torch.ops.aten.add.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -62,7 +62,7 @@ def test_add_scalar(self):

def test_addmm(self):
args = (torch.rand((5)), torch.rand((10, 5)), torch.rand((5, 5)))
dynamic_shapes = ((None, {0: Dim("dim")}, None),)
dynamic_shapes = (None, {0: Dim("dim")}, None)
m = wrap_func_as_nn_module(torch.ops.aten.addmm.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -83,7 +83,7 @@ def test_bmm(self):
torch.rand((24, 197, 64)),
torch.rand((24, 64, 197)),
)
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.bmm.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -104,7 +104,7 @@ def test_bmm_dynamic_out_dim(self):
torch.rand((8, 128, 256)),
torch.rand((8, 256, 3)),
)
dynamic_shapes = ((None, {2: Dim("dim")}),)
dynamic_shapes = (None, {2: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.bmm.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -125,7 +125,7 @@ def test_bmm_dynamic_reduction_dim(self):
torch.rand((8, 128, 3)),
torch.rand((8, 3, 256)),
)
dynamic_shapes = (({2: Dim("dim")}, {1: Dim("dim")}),)
dynamic_shapes = ({2: Dim("dim")}, {1: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.bmm.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -143,7 +143,7 @@ def test_bmm_dynamic_reduction_dim(self):

def test_cat(self):
args = (torch.rand((10, 1, 768)), torch.rand((10, 196, 768)))
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(
lambda x, y: torch.ops.aten.cat.default([x, y], 1))
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
Expand All @@ -166,7 +166,7 @@ def test_conv(self):
torch.rand((5, 3, 16, 16)),
torch.rand((5)),
)
dynamic_shapes = (({0: Dim("dim")}, None, None),)
dynamic_shapes = ({0: Dim("dim")}, None, None)
m = wrap_func_as_nn_module(
lambda x, y, z: torch.ops.aten.convolution.default(
x,
Expand Down Expand Up @@ -197,7 +197,7 @@ def test_conv1d(self):
torch.rand((3, 1, 800)),
torch.rand((512, 1, 10)),
)
dynamic_shapes = (({0: Dim("dim")}, None),)
dynamic_shapes = ({0: Dim("dim")}, None)
# dynamic_shapes = None
m = wrap_func_as_nn_module(lambda x, y: torch.ops.aten.convolution.default(
x,
Expand Down Expand Up @@ -225,7 +225,7 @@ def test_conv1d(self):

def test_cumsum(self):
args = (torch.rand((10, 5)), 1)
dynamic_shapes = (({0: Dim("dim")}, None),)
dynamic_shapes = ({0: Dim("dim")}, None)
m = wrap_func_as_nn_module(torch.ops.aten.cumsum.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -242,7 +242,7 @@ def test_cumsum(self):

def test_div(self):
args = (torch.rand((10, 12, 197)), torch.rand((10, 12, 197)))
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.div.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -260,7 +260,7 @@ def test_div(self):

def test_div_scalar(self):
args = (torch.rand((10, 12, 197)), 8.0)
dynamic_shapes = (({0: Dim("dim")}, None),)
dynamic_shapes = ({0: Dim("dim")}, None)
m = wrap_func_as_nn_module(torch.ops.aten.div.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -277,7 +277,7 @@ def test_div_scalar(self):

def test_gelu(self):
args = (torch.rand((3, 5)),)
dynamic_shapes = (({0: Dim("dim")},),)
dynamic_shapes = ({0: Dim("dim")},)
m = wrap_func_as_nn_module(torch.ops.aten.gelu)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand Down Expand Up @@ -342,7 +342,7 @@ def forward(self, x):

def test_mul(self):
args = (torch.rand((10, 2, 768)), torch.rand((10, 2, 768)))
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.mul.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -360,7 +360,7 @@ def test_mul(self):

def test_mul_scalar(self):
args = (torch.rand((10, 2, 768)), 0.125)
dynamic_shapes = (({0: Dim("dim")}, None),)
dynamic_shapes = ({0: Dim("dim")}, None)
m = wrap_func_as_nn_module(torch.ops.aten.mul.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand Down Expand Up @@ -483,7 +483,7 @@ def forward(self, x, weight, bias):

def test_permute(self):
args = (torch.rand((10, 197, 12, 64)),)
dynamic_shapes = (({0: Dim("dim")},),)
dynamic_shapes = ({0: Dim("dim")},)
m = wrap_func_as_nn_module(
lambda x: torch.ops.aten.permute.default(x, [0, 2, 1, 3]))
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
Expand All @@ -502,7 +502,7 @@ def test_permute(self):

def test_select(self):
args = (torch.rand((10, 197, 768)), 1, 0)
dynamic_shapes = (({0: Dim("dim")}, None, None),)
dynamic_shapes = ({0: Dim("dim")}, None, None)
m = wrap_func_as_nn_module(torch.ops.aten.select.int)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -519,7 +519,7 @@ def test_select(self):

def test_slice(self):
args = (torch.rand((10, 3, 224, 224)), 0, 0, 9223372036854775807)
dynamic_shapes = (({0: Dim("dim")}, None, None, None),)
dynamic_shapes = ({0: Dim("dim")}, None, None, None)
m = wrap_func_as_nn_module(torch.ops.aten.slice.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -537,7 +537,7 @@ def test_slice(self):

def test_slice_2(self):
args = (torch.rand((10, 3, 224, 224)), 1, 0, 2)
dynamic_shapes = (({0: Dim("dim")}, None, None, None),)
dynamic_shapes = ({0: Dim("dim")}, None, None, None)
m = wrap_func_as_nn_module(torch.ops.aten.slice.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -555,7 +555,7 @@ def test_slice_2(self):

def test_softmax(self):
args = (torch.rand((10, 12, 197, 197)), -1, False)
dynamic_shapes = (({0: Dim("dim")}, None, None),)
dynamic_shapes = ({0: Dim("dim")}, None, None)
m = wrap_func_as_nn_module(torch.ops.aten._softmax.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -573,7 +573,7 @@ def test_softmax(self):

def test_sub(self):
args = (torch.rand((10, 1, 1, 10)), torch.rand((10, 1, 1, 10)))
dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)
dynamic_shapes = ({0: Dim("dim")}, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.sub.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -591,7 +591,7 @@ def test_sub(self):

def test_softmax_reduce_on_dynamic_dim(self):
args = (torch.rand((1, 8, 128, 3)), -1, False)
dynamic_shapes = (({3: Dim("dim")}, None, None),)
dynamic_shapes = ({3: Dim("dim")}, None, None)
m = wrap_func_as_nn_module(torch.ops.aten._softmax.default)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand Down Expand Up @@ -628,7 +628,7 @@ def test_index(self):

def test_sub_scalar(self):
args = (1.0, torch.rand((10, 1, 1, 10)))
dynamic_shapes = ((None, {0: Dim("dim")}),)
dynamic_shapes = (None, {0: Dim("dim")})
m = wrap_func_as_nn_module(torch.ops.aten.sub.Tensor)
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand Down Expand Up @@ -670,7 +670,7 @@ def forward(self, x):

def test_transpose_on_dynamic_dim(self):
args = (torch.rand((1, 8, 3, 256)),)
dynamic_shapes = (({2: Dim("dim")},),)
dynamic_shapes = ({2: Dim("dim")},)
m = wrap_func_as_nn_module(
lambda x: torch.ops.aten.transpose.int(x, -2, -1))
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
Expand All @@ -688,7 +688,7 @@ def test_transpose_on_dynamic_dim(self):

def test_unsqueeze_1(self):
args = (torch.rand((3, 10)),)
dynamic_shapes = (({0: Dim("dim")},),)
dynamic_shapes = ({0: Dim("dim")},)
m = wrap_func_as_nn_module(lambda x: torch.ops.aten.unsqueeze.default(x, 1))
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand All @@ -705,7 +705,7 @@ def test_unsqueeze_1(self):

def test_unsqueeze_2(self):
args = (torch.rand((1, 1, 3, 256)),)
dynamic_shapes = (({2: Dim("dim")},),)
dynamic_shapes = ({2: Dim("dim")},)
m = wrap_func_as_nn_module(lambda x: torch.ops.aten.unsqueeze.default(x, 2))
ep = export(m, args=args, dynamic_shapes=dynamic_shapes)
shlo_module = exported_program_to_stablehlo(ep)
Expand Down
Loading