Skip to content

Commit

Permalink
add more test cases
Browse files Browse the repository at this point in the history
  • Loading branch information
zewenli98 committed Apr 26, 2024
1 parent ade7ed6 commit fbed329
Showing 1 changed file with 19 additions and 6 deletions.
25 changes: 19 additions & 6 deletions tests/py/dynamo/conversion/test_layer_norm_aten.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,18 +39,31 @@ def forward(self, x):


class TestNativeLayerNormConverter(DispatchTestCase):
def test_layer_norm(self):
@parameterized.expand(
[
(
(5, 3, 2, 4),
[
4,
],
),
((5, 3, 2, 4), [2, 4]),
((5, 3, 2, 4), [3, 2, 4]),
((5, 3, 2, 4), [5, 3, 2, 4]),
]
)
def test_layer_norm(self, input_shape, normalized_shape, eps=1e-05):
class LayerNorm(torch.nn.Module):
def forward(self, x):
return torch.ops.aten.native_layer_norm.default(
x,
torch.tensor([3, 224, 224]),
torch.ones((3, 224, 224)),
torch.zeros((3, 224, 224)),
1e-05,
normalized_shape,
torch.randn(normalized_shape),
torch.randn(normalized_shape),
eps,
)[0]

inputs = [torch.randn(1, 3, 224, 224)]
inputs = [torch.randn(input_shape)]
self.run_test(
LayerNorm(),
inputs,
Expand Down

0 comments on commit fbed329

Please sign in to comment.