[Core Aten] Add and enable tests for aten index_select and logical_and (

pytorch#6293)
ManfeiBai · Jan 11, 2024 · 04e1238 · 04e1238
1 parent 8141078
commit 04e1238
Showing 1 changed file with 35 additions and 2 deletions.
diff --git a/test/test_core_aten_ops.py b/test/test_core_aten_ops.py
@@ -2162,7 +2162,6 @@ def test_aten_index_select_0(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.index_select, args, kwargs)
 
-  @unittest.skip
   def test_aten_index_select_1(self):
     args = (
         torch.randn((2, 10)).to(torch.float16),
@@ -2181,6 +2180,33 @@ def test_aten_index_select_2(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.index_select, args, kwargs)
 
+  def test_aten_index_select_3(self):
+    args = (
+        torch.randn((2, 10)).to(torch.float32),
+        1,
+        torch.randint(0, 10, (2,)).to(torch.int32),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.index_select, args, kwargs)
+
+  def test_aten_index_select_4(self):
+    args = (
+        torch.randn((2, 10)).to(torch.float16),
+        1,
+        torch.randint(0, 10, (2,)).to(torch.int32),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.index_select, args, kwargs)
+
+  def test_aten_index_select_5(self):
+    args = (
+        torch.randint(0, 10, (2, 10)).to(torch.int32),
+        1,
+        torch.randint(0, 10, (2,)).to(torch.int32),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.index_select, args, kwargs)
+
   @unittest.skip
   def test_aten_index_Tensor_0(self):
     args = (
@@ -2437,7 +2463,6 @@ def test_aten_logical_and_1(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.logical_and, args, kwargs)
 
-  @unittest.skip
   def test_aten_logical_and_2(self):
     args = (
         torch.randint(0, 10, (10, 10)).to(torch.int32),
@@ -2446,6 +2471,14 @@ def test_aten_logical_and_2(self):
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.logical_and, args, kwargs)
 
+  def test_aten_logical_and_3(self):
+    args = (
+        torch.randint(0, 2, (10, 10)).to(torch.bool),
+        torch.randint(0, 2, (10, 10)).to(torch.bool),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.logical_and, args, kwargs)
+
   def test_aten_logical_not_0(self):
     args = (torch.randn((10, 10)).to(torch.float32),)
     kwargs = dict()