diff --git a/backends/test/suite/operators/test_masked_fill.py b/backends/test/suite/operators/test_masked_fill.py new file mode 100644 index 00000000000..68dccba69f3 --- /dev/null +++ b/backends/test/suite/operators/test_masked_fill.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import Union + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class MaskedFillModel(torch.nn.Module): + def __init__(self, value: Union[float, int]): + super().__init__() + self.value = value + + def forward(self, x, mask): + return x.masked_fill(mask, self.value) + + +@operator_test +class MaskedFill(OperatorTest): + @dtype_test + def test_masked_fill_dtype(self, flow: TestFlow, dtype) -> None: + mask = torch.randint(0, 2, (16, 32), dtype=torch.bool) + self._test_op( + MaskedFillModel(value=0.0), + ( + torch.rand(16, 32).to(dtype), + mask, + ), + flow, + ) + + def test_masked_fill_different_values(self, flow: TestFlow) -> None: + mask = torch.randint(0, 2, (16, 32), dtype=torch.bool) + + self._test_op( + MaskedFillModel(value=5.0), + ( + torch.randn(16, 32), + mask, + ), + flow, + ) + + self._test_op( + MaskedFillModel(value=-5.0), + ( + torch.randn(16, 32), + mask, + ), + flow, + ) + + self._test_op( + MaskedFillModel(value=1), + ( + torch.randn(16, 32), + mask, + ), + flow, + ) + + def test_masked_fill_different_shapes(self, flow: TestFlow) -> None: + self._test_op( + MaskedFillModel(value=0.0), + ( + torch.randn(512), + torch.randint(0, 2, (512,), dtype=torch.bool), + ), + flow, + ) + + self._test_op( + MaskedFillModel(value=0.0), + ( + torch.randn(4, 8, 16), + torch.randint(0, 2, (4, 8, 16), dtype=torch.bool), + ), + flow, + ) + + def test_masked_fill_broadcast(self, flow: TestFlow) -> None: + self._test_op( + MaskedFillModel(value=0.0), + ( + torch.randn(16, 32), + torch.randint(0, 2, (32,), dtype=torch.bool), + ), + flow, + ) diff --git a/backends/test/suite/operators/test_permute.py b/backends/test/suite/operators/test_permute.py new file mode 100644 index 00000000000..bc79a63d610 --- /dev/null +++ b/backends/test/suite/operators/test_permute.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import List + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class PermuteModel(torch.nn.Module): + def __init__(self, dims: List[int]): + super().__init__() + self.dims = dims + + def forward(self, x): + return x.permute(self.dims) + + +@operator_test +class Permute(OperatorTest): + @dtype_test + def test_permute_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + PermuteModel(dims=[1, 0]), + (torch.rand(20, 32).to(dtype),), + flow, + ) + + def test_permute_3d(self, flow: TestFlow) -> None: + self._test_op( + PermuteModel(dims=[2, 0, 1]), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + PermuteModel(dims=[1, 2, 0]), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + PermuteModel(dims=[0, 2, 1]), + (torch.randn(8, 10, 12),), + flow, + ) + + def test_permute_4d(self, flow: TestFlow) -> None: + self._test_op( + PermuteModel(dims=[3, 2, 1, 0]), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + self._test_op( + PermuteModel(dims=[0, 2, 1, 3]), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + def test_permute_identity(self, flow: TestFlow) -> None: + self._test_op( + PermuteModel(dims=[0, 1]), + (torch.randn(20, 32),), + flow, + ) + + self._test_op( + PermuteModel(dims=[0, 1, 2]), + (torch.randn(8, 10, 12),), + flow, + ) + + def test_permute_negative_dims(self, flow: TestFlow) -> None: + self._test_op( + PermuteModel(dims=[-1, -3, -2, -4]), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + self._test_op( + PermuteModel(dims=[-4, -2, -3, -1]), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + def test_permute_different_shapes(self, flow: TestFlow) -> None: + self._test_op( + PermuteModel(dims=[0]), + (torch.randn(512),), + flow, + ) + + self._test_op( + PermuteModel(dims=[4, 3, 2, 1, 0]), + (torch.randn(2, 3, 4, 5, 6),), + flow, + ) diff --git a/backends/test/suite/operators/test_transpose.py b/backends/test/suite/operators/test_transpose.py new file mode 100644 index 00000000000..77f5e62cb18 --- /dev/null +++ b/backends/test/suite/operators/test_transpose.py @@ -0,0 +1,143 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class TransposeModel(torch.nn.Module): + def __init__(self, dim0: int, dim1: int): + super().__init__() + self.dim0 = dim0 + self.dim1 = dim1 + + def forward(self, x): + return torch.transpose(x, self.dim0, self.dim1) + + +@operator_test +class Transpose(OperatorTest): + @dtype_test + def test_transpose_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=1), + (torch.rand(20, 32).to(dtype),), + flow, + ) + + def test_transpose_basic(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=1), + (torch.randn(20, 32),), + flow, + ) + + def test_transpose_3d(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=1), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + TransposeModel(dim0=0, dim1=2), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + TransposeModel(dim0=1, dim1=2), + (torch.randn(8, 10, 12),), + flow, + ) + + def test_transpose_4d(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=3), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + self._test_op( + TransposeModel(dim0=1, dim1=2), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + def test_transpose_identity(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=0), + (torch.randn(20, 32),), + flow, + ) + self._test_op( + TransposeModel(dim0=1, dim1=1), + (torch.randn(20, 32),), + flow, + ) + + self._test_op( + TransposeModel(dim0=0, dim1=0), + (torch.randn(8, 10, 12),), + flow, + ) + self._test_op( + TransposeModel(dim0=1, dim1=1), + (torch.randn(8, 10, 12),), + flow, + ) + self._test_op( + TransposeModel(dim0=2, dim1=2), + (torch.randn(8, 10, 12),), + flow, + ) + + def test_transpose_negative_dims(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=-3, dim1=-1), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + TransposeModel(dim0=-2, dim1=-1), + (torch.randn(8, 10, 12),), + flow, + ) + + def test_transpose_different_shapes(self, flow: TestFlow) -> None: + self._test_op( + TransposeModel(dim0=0, dim1=1), + (torch.randn(20, 32),), + flow, + ) + + self._test_op( + TransposeModel(dim0=0, dim1=2), + (torch.randn(8, 10, 12),), + flow, + ) + + self._test_op( + TransposeModel(dim0=1, dim1=3), + (torch.randn(4, 6, 8, 10),), + flow, + ) + + self._test_op( + TransposeModel(dim0=0, dim1=4), + (torch.randn(2, 3, 4, 5, 6),), + flow, + )