[Backend Tester] Add index_put and index_select tests

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)

backends/test/suite/operators/test_cat.py

@@ -0,0 +1,176 @@
+# 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 CatModel(torch.nn.Module):
+    def __init__(self, dim: int = 0):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x1, x2, x3):
+        return torch.cat([x1, x2, x3], dim=self.dim)
+
+
+@operator_test
+class Cat(OperatorTest):
+    @dtype_test
+    def test_cat_dtype(self, flow: TestFlow, dtype) -> None:
+        self._test_op(
+            CatModel(),
+            (
+                torch.rand(8, 32).to(dtype),
+                torch.rand(12, 32).to(dtype),
+                torch.rand(16, 32).to(dtype),
+            ),
+            flow,
+        )
+
+    def test_cat_dimensions(self, flow: TestFlow) -> None:
+        self._test_op(
+            CatModel(dim=0),
+            (
+                torch.randn(8, 32),
+                torch.randn(12, 32),
+                torch.randn(16, 32),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=1),
+            (
+                torch.randn(16, 8),
+                torch.randn(16, 12),
+                torch.randn(16, 16),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=2),
+            (
+                torch.randn(4, 8, 4),
+                torch.randn(4, 8, 8),
+                torch.randn(4, 8, 12),
+            ),
+            flow,
+        )
+
+    def test_cat_negative_dim(self, flow: TestFlow) -> None:
+        self._test_op(
+            CatModel(dim=-1),
+            (
+                torch.randn(16, 8),
+                torch.randn(16, 12),
+                torch.randn(16, 16),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=-2),
+            (
+                torch.randn(8, 32),
+                torch.randn(12, 32),
+                torch.randn(16, 32),
+            ),
+            flow,
+        )
+
+    def test_cat_different_shapes(self, flow: TestFlow) -> None:
+        self._test_op(
+            CatModel(),
+            (
+                torch.randn(128),
+                torch.randn(256),
+                torch.randn(384),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=0),
+            (
+                torch.randn(4, 8, 16),
+                torch.randn(8, 8, 16),
+                torch.randn(12, 8, 16),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=1),
+            (
+                torch.randn(8, 4, 16),
+                torch.randn(8, 8, 16),
+                torch.randn(8, 12, 16),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=2),
+            (
+                torch.randn(8, 12, 4),
+                torch.randn(8, 12, 8),
+                torch.randn(8, 12, 12),
+            ),
+            flow,
+        )
+
+    def test_cat_broadcast(self, flow: TestFlow) -> None:
+        self._test_op(
+            CatModel(dim=0),
+            (
+                torch.randn(2, 16, 32),
+                torch.randn(4, 16, 32),
+                torch.randn(6, 16, 32),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=1),
+            (
+                torch.randn(8, 8, 16),
+                torch.randn(8, 16, 16),
+                torch.randn(8, 24, 16),
+            ),
+            flow,
+        )
+
+        self._test_op(
+            CatModel(dim=2),
+            (
+                torch.randn(4, 16, 8),
+                torch.randn(4, 16, 16),
+                torch.randn(4, 16, 24),
+            ),
+            flow,
+        )
+
+    def test_cat_same_shapes(self, flow: TestFlow) -> None:
+        self._test_op(
+            CatModel(),
+            (
+                torch.randn(8, 32),
+                torch.randn(8, 32),
+                torch.randn(8, 32),
+            ),
+            flow,
+        )

backends/test/suite/operators/test_expand.py

@@ -0,0 +1,122 @@
+# 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 ExpandModel(torch.nn.Module):
+    def __init__(self, shape: List[int]):
+        super().__init__()
+        self.shape = shape
+
+    def forward(self, x):
+        return x.expand(self.shape)
+
+
+@operator_test
+class Expand(OperatorTest):
+    @dtype_test
+    def test_expand_dtype(self, flow: TestFlow, dtype) -> None:
+        self._test_op(
+            ExpandModel(shape=[8, 32]),
+            (torch.rand(1, 32).to(dtype),),
+            flow,
+        )
+
+    def test_expand_dimensions(self, flow: TestFlow) -> None:
+        self._test_op(
+            ExpandModel(shape=[8, 32]),
+            (torch.randn(1, 32),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[16, 20]),
+            (torch.randn(1, 1),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[4, 1, 32]),
+            (torch.randn(1, 32),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[8, 4, 16]),
+            (torch.randn(8, 1, 16),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[6, 16, 8]),
+            (torch.randn(6, 16, 1),),
+            flow,
+        )
+
+    def test_expand_keep_original_size(self, flow: TestFlow) -> None:
+        self._test_op(
+            ExpandModel(shape=[8, -1]),
+            (torch.randn(1, 32),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[-1, 32]),
+            (torch.randn(4, 1),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[-1, 16, -1]),
+            (torch.randn(4, 1, 8),),
+            flow,
+        )
+
+    def test_expand_rank_increase(self, flow: TestFlow) -> None:
+        # Test expanding 2D tensor to 3D
+        self._test_op(
+            ExpandModel(shape=[6, 8, 16]),
+            (torch.randn(8, 16),),
+            flow,
+        )
+
+        # Test expanding 2D tensor to 4D
+        self._test_op(
+            ExpandModel(shape=[3, 4, 8, 16]),
+            (torch.randn(8, 16),),
+            flow,
+        )
+
+    def test_expand_singleton_dimensions(self, flow: TestFlow) -> None:
+        self._test_op(
+            ExpandModel(shape=[512]),
+            (torch.randn(1),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[16, 20]),
+            (torch.randn(1, 1),),
+            flow,
+        )
+
+        self._test_op(
+            ExpandModel(shape=[8, 32]),
+            (torch.randn(32),),
+            flow,
+        )