diff --git a/backends/test/harness/tester.py b/backends/test/harness/tester.py index 06db1aae13d..7019b734290 100644 --- a/backends/test/harness/tester.py +++ b/backends/test/harness/tester.py @@ -311,7 +311,10 @@ def run_method_and_compare_outputs( print(f"Comparing Stage {stage} with Stage {reference_stage}") for run_iteration in range(number_of_runs): inputs_to_run = inputs if inputs else next(self.generate_random_inputs()) - input_shapes = [generated_input.shape for generated_input in inputs_to_run] + input_shapes = [ + generated_input.shape if hasattr(generated_input, "shape") else None + for generated_input in inputs_to_run + ] print(f"Run {run_iteration} with input shapes: {input_shapes}") # Reference output (and quantization scale) diff --git a/backends/test/suite/operators/test_amax.py b/backends/test/suite/operators/test_amax.py new file mode 100644 index 00000000000..aff33476e69 --- /dev/null +++ b/backends/test/suite/operators/test_amax.py @@ -0,0 +1,255 @@ +# 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, Optional, Tuple, Union + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class AmaxModel(torch.nn.Module): + def __init__( + self, + dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None, + keepdim: bool = False, + ): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.amax(x, dim=self.dim, keepdim=self.keepdim) + + +@operator_test +class Amax(OperatorTest): + @dtype_test + def test_amax_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + AmaxModel().to(dtype), + (torch.rand(10, 10).to(dtype),), + flow, + ) + + def test_amax_dim(self, flow: TestFlow) -> None: + self._test_op( + AmaxModel(dim=0), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AmaxModel(dim=1), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AmaxModel(dim=0), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=2), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=1), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=-1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=-2), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_amax_multi_dim(self, flow: TestFlow) -> None: + self._test_op( + AmaxModel(dim=(0, 1)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(0, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(1, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(1, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(0, 2)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(-1, -3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(0, 1, 2, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + def test_amax_keepdim(self, flow: TestFlow) -> None: + self._test_op( + AmaxModel(dim=0, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AmaxModel(dim=1, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AmaxModel(dim=1, keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=2, keepdim=True), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(dim=(1, 2), keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_amax_shapes(self, flow: TestFlow) -> None: + self._test_op( + AmaxModel(), + (torch.randn(20),), + flow, + ) + self._test_op( + AmaxModel(dim=0), + (torch.randn(20),), + flow, + ) + + self._test_op( + AmaxModel(), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AmaxModel(), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AmaxModel(), + (torch.randn(2, 2, 3, 4, 5),), + flow, + ) + + def test_amax_edge_cases(self, flow: TestFlow) -> None: + x = torch.tensor([[1.0, float("inf"), 3.0], [4.0, 5.0, float("inf")]]) + self._test_op( + AmaxModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AmaxModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AmaxModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + AmaxModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AmaxModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AmaxModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + def test_amax_scalar(self, flow: TestFlow) -> None: + self._test_op( + AmaxModel(), + (torch.tensor([5.0]),), + flow, + ) + self._test_op( + AmaxModel(dim=0), + (torch.tensor([5.0]),), + flow, + ) diff --git a/backends/test/suite/operators/test_amin.py b/backends/test/suite/operators/test_amin.py new file mode 100644 index 00000000000..ab59d77d0be --- /dev/null +++ b/backends/test/suite/operators/test_amin.py @@ -0,0 +1,257 @@ +# 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, Optional, Tuple, Union + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class AminModel(torch.nn.Module): + def __init__( + self, + dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None, + keepdim: bool = False, + ): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + if self.dim is None: + return torch.amin(x, keepdim=self.keepdim) + return torch.amin(x, dim=self.dim, keepdim=self.keepdim) + + +@operator_test +class Amin(OperatorTest): + @dtype_test + def test_amin_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + AminModel().to(dtype), + (torch.rand(10, 10).to(dtype),), + flow, + ) + + def test_amin_dim(self, flow: TestFlow) -> None: + self._test_op( + AminModel(dim=0), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AminModel(dim=1), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AminModel(dim=0), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=2), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=1), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=-1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=-2), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_amin_multi_dim(self, flow: TestFlow) -> None: + self._test_op( + AminModel(dim=(0, 1)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(0, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(1, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(1, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(0, 2)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(-1, -3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(0, 1, 2, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + def test_amin_keepdim(self, flow: TestFlow) -> None: + self._test_op( + AminModel(dim=0, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AminModel(dim=1, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AminModel(dim=1, keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=2, keepdim=True), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(dim=(1, 2), keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_amin_shapes(self, flow: TestFlow) -> None: + self._test_op( + AminModel(), + (torch.randn(20),), + flow, + ) + self._test_op( + AminModel(dim=0), + (torch.randn(20),), + flow, + ) + + self._test_op( + AminModel(), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + AminModel(), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + AminModel(), + (torch.randn(2, 2, 3, 4, 5),), + flow, + ) + + def test_amin_edge_cases(self, flow: TestFlow) -> None: + x = torch.tensor([[1.0, float("-inf"), 3.0], [4.0, 5.0, float("-inf")]]) + self._test_op( + AminModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AminModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AminModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + AminModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AminModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + AminModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + def test_amin_scalar(self, flow: TestFlow) -> None: + self._test_op( + AminModel(), + (torch.tensor([5.0]),), + flow, + ) + self._test_op( + AminModel(dim=0), + (torch.tensor([5.0]),), + flow, + ) diff --git a/backends/test/suite/operators/test_argmax.py b/backends/test/suite/operators/test_argmax.py new file mode 100644 index 00000000000..adf1e43a340 --- /dev/null +++ b/backends/test/suite/operators/test_argmax.py @@ -0,0 +1,199 @@ +# 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 Optional + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class ArgmaxModel(torch.nn.Module): + def __init__(self, dim: Optional[int] = None, keepdim: bool = False): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.argmax(x, dim=self.dim, keepdim=self.keepdim) + + +@operator_test +class Argmax(OperatorTest): + @dtype_test + def test_argmax_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + ArgmaxModel().to(dtype), + (torch.rand(10, 10).to(dtype),), + flow, + ) + + def test_argmax_dim(self, flow: TestFlow) -> None: + self._test_op( + ArgmaxModel(dim=0), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=1), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=0), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=2), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=1), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=-1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=-2), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_argmax_keepdim(self, flow: TestFlow) -> None: + self._test_op( + ArgmaxModel(dim=0, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=1, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=1, keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(dim=2, keepdim=True), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + def test_argmax_shapes(self, flow: TestFlow) -> None: + self._test_op( + ArgmaxModel(), + (torch.randn(20),), + flow, + ) + + self._test_op( + ArgmaxModel(), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgmaxModel(), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + ArgmaxModel(), + (torch.randn(2, 2, 3, 4, 5),), + flow, + ) + + def test_argmax_edge_cases(self, flow: TestFlow) -> None: + x = torch.tensor([[1.0, float("inf"), 3.0], [4.0, 5.0, float("inf")]]) + self._test_op( + ArgmaxModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgmaxModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgmaxModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + ArgmaxModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgmaxModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgmaxModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([5.0]) + self._test_op( + ArgmaxModel(), + (x,), + flow, + ) + + def test_argmax_scalar(self, flow: TestFlow) -> None: + self._test_op( + ArgmaxModel(), + (torch.tensor([5.0]),), + flow, + ) diff --git a/backends/test/suite/operators/test_argmin.py b/backends/test/suite/operators/test_argmin.py new file mode 100644 index 00000000000..0613c74a3ee --- /dev/null +++ b/backends/test/suite/operators/test_argmin.py @@ -0,0 +1,199 @@ +# 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 Optional + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class ArgminModel(torch.nn.Module): + def __init__(self, dim: Optional[int] = None, keepdim: bool = False): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.argmin(x, dim=self.dim, keepdim=self.keepdim) + + +@operator_test +class Argmin(OperatorTest): + @dtype_test + def test_argmin_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + ArgminModel().to(dtype), + (torch.rand(10, 10).to(dtype),), + flow, + ) + + def test_argmin_dim(self, flow: TestFlow) -> None: + self._test_op( + ArgminModel(dim=0), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgminModel(dim=1), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgminModel(dim=0), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=2), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=1), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=-1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=-2), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_argmin_keepdim(self, flow: TestFlow) -> None: + self._test_op( + ArgminModel(dim=0, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgminModel(dim=1, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgminModel(dim=1, keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(dim=2, keepdim=True), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + def test_argmin_shapes(self, flow: TestFlow) -> None: + self._test_op( + ArgminModel(), + (torch.randn(20),), + flow, + ) + + self._test_op( + ArgminModel(), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + ArgminModel(), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + ArgminModel(), + (torch.randn(2, 2, 3, 4, 5),), + flow, + ) + + def test_argmin_edge_cases(self, flow: TestFlow) -> None: + x = torch.tensor([[1.0, float("-inf"), 3.0], [4.0, 5.0, float("-inf")]]) + self._test_op( + ArgminModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgminModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgminModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + ArgminModel(), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgminModel(dim=0), + (x,), + flow, + use_random_test_inputs=False, + ) + self._test_op( + ArgminModel(dim=1), + (x,), + flow, + use_random_test_inputs=False, + ) + + x = torch.tensor([5.0]) + self._test_op( + ArgminModel(), + (x,), + flow, + ) + + def test_argmin_scalar(self, flow: TestFlow) -> None: + self._test_op( + ArgminModel(), + (torch.tensor([5.0]),), + flow, + ) diff --git a/backends/test/suite/operators/test_index_put.py b/backends/test/suite/operators/test_index_put.py new file mode 100644 index 00000000000..b5333b40984 --- /dev/null +++ b/backends/test/suite/operators/test_index_put.py @@ -0,0 +1,455 @@ +# 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 IndexPutInPlaceModel(torch.nn.Module): + def __init__(self, accumulate=False): + super().__init__() + self.accumulate = accumulate + + def forward(self, x, indices, values): + # Clone the input to avoid modifying it in-place + result = x.clone() + # Apply index_put_ and return the modified tensor + result.index_put_(indices, values, self.accumulate) + return result + + +class IndexPutModel(torch.nn.Module): + def __init__(self, accumulate=False): + super().__init__() + self.accumulate = accumulate + + def forward(self, x, indices, values): + # Use the non-in-place variant which returns a new tensor + return torch.index_put(x, indices, values, self.accumulate) + + +@operator_test +class IndexPut(OperatorTest): + @dtype_test + def test_index_put_in_place_dtype(self, flow: TestFlow, dtype) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]).to(dtype) + self._test_op( + IndexPutInPlaceModel(), + ((torch.rand(5, 2) * 100).to(dtype), indices, values), + flow, + generate_random_test_inputs=False, + ) + + @dtype_test + def test_index_put_dtype(self, flow: TestFlow, dtype) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]).to(dtype) + self._test_op( + IndexPutModel(), + ((torch.rand(5, 2) * 100).to(dtype), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_in_place_accumulate(self, flow: TestFlow) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=False), + (torch.ones(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=True), + (torch.ones(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_accumulate(self, flow: TestFlow) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutModel(accumulate=False), + (torch.ones(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutModel(accumulate=True), + (torch.ones(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_in_place_shapes(self, flow: TestFlow) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]), torch.tensor([1, 1])) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]), torch.tensor([1, 1]), torch.tensor([0, 1])) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = ( + torch.tensor([0, 2]), + torch.tensor([1, 1]), + torch.tensor([0, 1]), + torch.tensor([2, 3]), + ) + values = torch.tensor( + [ + 10.0, + ] + ) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 3, 2, 4), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_shapes(self, flow: TestFlow) -> None: + indices = (torch.tensor([0, 2]),) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]), torch.tensor([1, 1])) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2]), torch.tensor([1, 1]), torch.tensor([0, 1])) + values = torch.tensor([10.0, 20.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5, 3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = ( + torch.tensor([0, 2]), + torch.tensor([1, 1]), + torch.tensor([0, 1]), + torch.tensor([2, 3]), + ) + values = torch.tensor( + [ + 10.0, + ] + ) + self._test_op( + IndexPutModel(), + (torch.randn(5, 3, 2, 4), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_in_place_indices(self, flow: TestFlow) -> None: + indices = (torch.tensor([2]),) + values = torch.tensor([10.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2, 4]),) + values = torch.tensor([10.0, 20.0, 30.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([1, 1, 3, 3]),) + values = torch.tensor([10.0, 20.0, 30.0, 40.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=True), + (torch.randn(5), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_indices(self, flow: TestFlow) -> None: + indices = (torch.tensor([2]),) + values = torch.tensor([10.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([0, 2, 4]),) + values = torch.tensor([10.0, 20.0, 30.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + indices = (torch.tensor([1, 1, 3, 3]),) + values = torch.tensor([10.0, 20.0, 30.0, 40.0]) + self._test_op( + IndexPutModel(accumulate=True), + (torch.randn(5), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_in_place_broadcasting(self, flow: TestFlow) -> None: + # Test scalar broadcasting - single value to multiple positions + indices = (torch.tensor([0, 2, 4]),) + values = torch.tensor([42.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(5, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test 1D broadcasting to 2D indexed positions + indices = (torch.tensor([0, 1]), torch.tensor([1, 2])) + values = torch.tensor([10.0, 20.0]) # 1D tensor + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(3, 4), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test broadcasting with compatible shapes - 1D to multiple 2D slices + indices = (torch.tensor([0, 2]),) + values = torch.tensor([5.0, 15.0]) # Will broadcast to (2, 3) shape + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(4, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test 2D values broadcasting to 3D indexed positions + indices = (torch.tensor([0, 1]),) + values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # 2D tensor + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(3, 2, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test broadcasting with accumulate=True + indices = (torch.tensor([1, 1, 1]),) + values = torch.tensor([5.0]) # Scalar will be added 3 times to same position + self._test_op( + IndexPutInPlaceModel(accumulate=True), + (torch.ones(4, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_broadcasting(self, flow: TestFlow) -> None: + # Test scalar broadcasting - single value to multiple positions + indices = (torch.tensor([0, 2, 4]),) + values = torch.tensor([42.0]) + self._test_op( + IndexPutModel(), + (torch.randn(5, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test 1D broadcasting to 2D indexed positions + indices = (torch.tensor([0, 1]), torch.tensor([1, 2])) + values = torch.tensor([10.0, 20.0]) # 1D tensor + self._test_op( + IndexPutModel(), + (torch.randn(3, 4), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test broadcasting with compatible shapes - 1D to multiple 2D slices + indices = (torch.tensor([0, 2]),) + values = torch.tensor([5.0, 15.0]) # Will broadcast to (2, 3) shape + self._test_op( + IndexPutModel(), + (torch.randn(4, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test 2D values broadcasting to 3D indexed positions + indices = (torch.tensor([0, 1]),) + values = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # 2D tensor + self._test_op( + IndexPutModel(), + (torch.randn(3, 2, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test broadcasting with accumulate=True + indices = (torch.tensor([1, 1, 1]),) + values = torch.tensor([5.0]) # Scalar will be added 3 times to same position + self._test_op( + IndexPutModel(accumulate=True), + (torch.ones(4, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_in_place_two_indices(self, flow: TestFlow) -> None: + # Test basic two-index tensor indexing + indices = (torch.tensor([0, 1, 2]), torch.tensor([1, 0, 2])) + values = torch.tensor([10.0, 20.0, 30.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(4, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with different lengths (broadcasting) + indices = (torch.tensor([0, 2]), torch.tensor([1, 1])) + values = torch.tensor([15.0, 25.0]) + self._test_op( + IndexPutInPlaceModel(), + (torch.randn(3, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with repeated positions and accumulate=True + indices = (torch.tensor([1, 1, 2]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=True), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with repeated positions and accumulate=False + indices = (torch.tensor([1, 1, 2]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=False), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with index broadcast. + indices = (torch.tensor([1]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutInPlaceModel(accumulate=False), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + def test_index_put_two_indices(self, flow: TestFlow) -> None: + # Test basic two-index tensor indexing + indices = (torch.tensor([0, 1, 2]), torch.tensor([1, 0, 2])) + values = torch.tensor([10.0, 20.0, 30.0]) + self._test_op( + IndexPutModel(), + (torch.randn(4, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with different lengths (broadcasting) + indices = (torch.tensor([0, 2]), torch.tensor([1, 1])) + values = torch.tensor([15.0, 25.0]) + self._test_op( + IndexPutModel(), + (torch.randn(3, 3), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with repeated positions and accumulate=True + indices = (torch.tensor([1, 1, 2]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutModel(accumulate=True), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with repeated positions and accumulate=False + indices = (torch.tensor([1, 1, 2]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutModel(accumulate=False), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) + + # Test two-index with index broadcast. + indices = (torch.tensor([1]), torch.tensor([0, 0, 1])) + values = torch.tensor([5.0, 10.0, 15.0]) + self._test_op( + IndexPutModel(accumulate=False), + (torch.zeros(3, 2), indices, values), + flow, + generate_random_test_inputs=False, + ) diff --git a/backends/test/suite/operators/test_index_select.py b/backends/test/suite/operators/test_index_select.py new file mode 100644 index 00000000000..46a8018ef93 --- /dev/null +++ b/backends/test/suite/operators/test_index_select.py @@ -0,0 +1,128 @@ +# 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 IndexSelectModel(torch.nn.Module): + def __init__(self, dim=0): + super().__init__() + self.dim = dim + + def forward(self, x, indices): + return torch.index_select(x, self.dim, indices) + + +@operator_test +class IndexSelect(OperatorTest): + @dtype_test + def test_index_select_dtype(self, flow: TestFlow, dtype) -> None: + indices = torch.tensor([0, 2], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + ((torch.rand(5, 3) * 100).to(dtype), indices), + flow, + generate_random_test_inputs=False, + ) + + def test_index_select_dimensions(self, flow: TestFlow) -> None: + indices = torch.tensor([0, 2], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + indices = torch.tensor([0, 1], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=1), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + indices = torch.tensor([0, 2], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=2), + (torch.randn(3, 4, 5), indices), + flow, + generate_random_test_inputs=False, + ) + + def test_index_select_shapes(self, flow: TestFlow) -> None: + indices = torch.tensor([0, 1], dtype=torch.int64) + + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5), indices), + flow, + generate_random_test_inputs=False, + ) + + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3, 2), indices), + flow, + generate_random_test_inputs=False, + ) + + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3, 2, 4), indices), + flow, + generate_random_test_inputs=False, + ) + + def test_index_select_indices(self, flow: TestFlow) -> None: + indices = torch.tensor([2], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + indices = torch.tensor([0, 2, 4], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + indices = torch.tensor([1, 1, 3, 3], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) + + indices = torch.tensor([4, 3, 2, 1, 0], dtype=torch.int64) + self._test_op( + IndexSelectModel(dim=0), + (torch.randn(5, 3), indices), + flow, + generate_random_test_inputs=False, + ) diff --git a/backends/test/suite/operators/test_mean.py b/backends/test/suite/operators/test_mean.py new file mode 100644 index 00000000000..746a4b16d9f --- /dev/null +++ b/backends/test/suite/operators/test_mean.py @@ -0,0 +1,303 @@ +# 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, Optional, Tuple, Union + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class MeanModel(torch.nn.Module): + def __init__( + self, + dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None, + keepdim: bool = False, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.dim = dim + self.keepdim = keepdim + self.dtype = dtype + + def forward(self, x): + return torch.mean(x, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype) + + +@operator_test +class Mean(OperatorTest): + @dtype_test + def test_mean_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + MeanModel().to(dtype), + (torch.rand(10, 10).to(dtype),), + flow, + ) + + def test_mean_basic(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(), + (torch.randn(10, 10),), + flow, + ) + + def test_mean_dim(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(dim=0), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(dim=1), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(dim=0), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=2), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=1), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=-1), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=-2), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_mean_multi_dim(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(dim=(0, 1)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(0, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(1, 2)), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(1, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(0, 2)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(-1, -3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(0, 1, 2, 3)), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + def test_mean_keepdim(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(dim=0, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(dim=1, keepdim=True), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(dim=1, keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=2, keepdim=True), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(dim=(1, 2), keepdim=True), + (torch.randn(3, 4, 5),), + flow, + ) + + def test_mean_output_dtype(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(dtype=torch.float32), + (torch.randint(0, 10, (5, 10)),), + flow, + ) + + self._test_op( + MeanModel(dtype=torch.float64), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(dim=1, dtype=torch.float64), + (torch.randn(5, 10),), + flow, + ) + + def test_mean_shapes(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(), + (torch.randn(20),), + flow, + ) + self._test_op( + MeanModel(dim=0), + (torch.randn(20),), + flow, + ) + + self._test_op( + MeanModel(), + (torch.randn(5, 10),), + flow, + ) + + self._test_op( + MeanModel(), + (torch.randn(3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(), + (torch.randn(2, 3, 4, 5),), + flow, + ) + + self._test_op( + MeanModel(), + (torch.randn(2, 2, 3, 4, 5),), + flow, + ) + + def test_mean_edge_cases(self, flow: TestFlow) -> None: + x = torch.tensor([[1.0, float("inf"), 3.0], [4.0, 5.0, float("inf")]]) + self._test_op( + MeanModel(), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=0), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=1), + (x,), + flow, + generate_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("-inf"), 3.0], [4.0, 5.0, float("-inf")]]) + self._test_op( + MeanModel(), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=0), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=1), + (x,), + flow, + generate_random_test_inputs=False, + ) + + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + MeanModel(), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=0), + (x,), + flow, + generate_random_test_inputs=False, + ) + self._test_op( + MeanModel(dim=1), + (x,), + flow, + generate_random_test_inputs=False, + ) + + def test_mean_scalar(self, flow: TestFlow) -> None: + self._test_op( + MeanModel(), + (torch.tensor([5.0]),), + flow, + ) + self._test_op( + MeanModel(dim=0), + (torch.tensor([5.0]),), + flow, + ) diff --git a/backends/test/suite/operators/test_median.py b/backends/test/suite/operators/test_median.py new file mode 100644 index 00000000000..93823b812ca --- /dev/null +++ b/backends/test/suite/operators/test_median.py @@ -0,0 +1,186 @@ +# 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 Optional + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class MedianModel(torch.nn.Module): + def __init__(self, dim: Optional[int] = None, keepdim: bool = False): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.median(x, dim=self.dim, keepdim=self.keepdim) + + +class MedianValueOnlyModel(torch.nn.Module): + """Model that returns only the median values (not indices) when dim is specified.""" + + def __init__(self, dim: Optional[int] = None, keepdim: bool = False): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + if self.dim is not None: + return torch.median(x, dim=self.dim, keepdim=self.keepdim)[0] + else: + return torch.median(x) + + +@operator_test +class Median(OperatorTest): + @dtype_test + def test_median_dtype(self, flow: TestFlow, dtype) -> None: + # Test with different dtypes (global reduction) + model = MedianValueOnlyModel().to(dtype) + self._test_op(model, (torch.rand(10, 10).to(dtype),), flow) + + def test_median_basic(self, flow: TestFlow) -> None: + # Basic test with default parameters (global reduction) + self._test_op(MedianValueOnlyModel(), (torch.randn(10, 10),), flow) + + def test_median_dim(self, flow: TestFlow) -> None: + # Test with different dimensions (values only) + + # 2D tensor, dim=0 + self._test_op(MedianValueOnlyModel(dim=0), (torch.randn(5, 10),), flow) + + # 2D tensor, dim=1 + self._test_op(MedianValueOnlyModel(dim=1), (torch.randn(5, 10),), flow) + + # 3D tensor, dim=0 + self._test_op(MedianValueOnlyModel(dim=0), (torch.randn(3, 4, 5),), flow) + + # 3D tensor, dim=1 + self._test_op(MedianValueOnlyModel(dim=1), (torch.randn(3, 4, 5),), flow) + + # 3D tensor, dim=2 + self._test_op(MedianValueOnlyModel(dim=2), (torch.randn(3, 4, 5),), flow) + + # 4D tensor, dim=1 + self._test_op(MedianValueOnlyModel(dim=1), (torch.randn(2, 3, 4, 5),), flow) + + # Negative dim (last dimension) + self._test_op(MedianValueOnlyModel(dim=-1), (torch.randn(3, 4, 5),), flow) + + # Negative dim (second-to-last dimension) + self._test_op(MedianValueOnlyModel(dim=-2), (torch.randn(3, 4, 5),), flow) + + def test_median_with_indices(self, flow: TestFlow) -> None: + # Test with different dimensions (values and indices) + + # 2D tensor, dim=0 + self._test_op(MedianModel(dim=0), (torch.randn(5, 10),), flow) + + # 2D tensor, dim=1 + self._test_op(MedianModel(dim=1), (torch.randn(5, 10),), flow) + + # 3D tensor, dim=0 + self._test_op(MedianModel(dim=0), (torch.randn(3, 4, 5),), flow) + + # 3D tensor, dim=1 + self._test_op(MedianModel(dim=1), (torch.randn(3, 4, 5),), flow) + + # 3D tensor, dim=2 + self._test_op(MedianModel(dim=2), (torch.randn(3, 4, 5),), flow) + + # 4D tensor, dim=1 + self._test_op(MedianModel(dim=1), (torch.randn(2, 3, 4, 5),), flow) + + # Negative dim (last dimension) + self._test_op(MedianModel(dim=-1), (torch.randn(3, 4, 5),), flow) + + # Negative dim (second-to-last dimension) + self._test_op(MedianModel(dim=-2), (torch.randn(3, 4, 5),), flow) + + def test_median_keepdim(self, flow: TestFlow) -> None: + # Test with keepdim=True (values only) + + # 2D tensor, dim=0, keepdim=True + self._test_op( + MedianValueOnlyModel(dim=0, keepdim=True), (torch.randn(5, 10),), flow + ) + + # 2D tensor, dim=1, keepdim=True + self._test_op( + MedianValueOnlyModel(dim=1, keepdim=True), (torch.randn(5, 10),), flow + ) + + # 3D tensor, dim=1, keepdim=True + self._test_op( + MedianValueOnlyModel(dim=1, keepdim=True), (torch.randn(3, 4, 5),), flow + ) + + # 4D tensor, dim=2, keepdim=True + self._test_op( + MedianValueOnlyModel(dim=2, keepdim=True), (torch.randn(2, 3, 4, 5),), flow + ) + + def test_median_keepdim_with_indices(self, flow: TestFlow) -> None: + # Test with keepdim=True (values and indices) + + # 2D tensor, dim=0, keepdim=True + self._test_op(MedianModel(dim=0, keepdim=True), (torch.randn(5, 10),), flow) + + # 2D tensor, dim=1, keepdim=True + self._test_op(MedianModel(dim=1, keepdim=True), (torch.randn(5, 10),), flow) + + # 3D tensor, dim=1, keepdim=True + self._test_op(MedianModel(dim=1, keepdim=True), (torch.randn(3, 4, 5),), flow) + + # 4D tensor, dim=2, keepdim=True + self._test_op( + MedianModel(dim=2, keepdim=True), (torch.randn(2, 3, 4, 5),), flow + ) + + def test_median_shapes(self, flow: TestFlow) -> None: + # Test with different tensor shapes (global reduction) + + # 1D tensor + self._test_op(MedianValueOnlyModel(), (torch.randn(20),), flow) + + # 2D tensor + self._test_op(MedianValueOnlyModel(), (torch.randn(5, 10),), flow) + + # 3D tensor + self._test_op(MedianValueOnlyModel(), (torch.randn(3, 4, 5),), flow) + + # 4D tensor + self._test_op(MedianValueOnlyModel(), (torch.randn(2, 3, 4, 5),), flow) + + # 5D tensor + self._test_op(MedianValueOnlyModel(), (torch.randn(2, 2, 3, 4, 5),), flow) + + def test_median_edge_cases(self, flow: TestFlow) -> None: + # Tensor with NaN (NaN should be propagated) + x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) + self._test_op( + MedianValueOnlyModel(), (x,), flow, generate_random_test_inputs=False + ) + self._test_op( + MedianValueOnlyModel(dim=0), (x,), flow, generate_random_test_inputs=False + ) + self._test_op( + MedianValueOnlyModel(dim=1), (x,), flow, generate_random_test_inputs=False + ) + + def test_median_scalar(self, flow: TestFlow) -> None: + # Test with scalar input (1-element tensor) + self._test_op(MedianValueOnlyModel(), (torch.tensor([5.0]),), flow) + self._test_op(MedianValueOnlyModel(dim=0), (torch.tensor([5.0]),), flow)