Skip associative_scan tests in ref eager mode (#433)

Author: yf225

test/test_associative_scan.py

@@ -99,6 +99,9 @@ def jit_add_combine_fn(x, y):
 
 
 class TestAssociativeScan(RefEagerTestBase, TestCase):
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_basic_addition(self):
         """Test basic associative_scan functionality with prefix sum."""
 
@@ -132,6 +135,9 @@ def test_scan_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertIn("param_0 + param_1", code)
         self.assertIn("tl.associative_scan", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_maximum(self):
         """Test associative_scan with maximum combine function."""
 
@@ -164,6 +170,9 @@ def test_max_kernel(x: torch.Tensor) -> torch.Tensor:
             "tl.maximum" in code or "triton_helpers.maximum" in code
         )
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_multiplication(self):
         """Test associative_scan with multiplication combine function."""
 
@@ -194,6 +203,9 @@ def test_mul_kernel(x: torch.Tensor) -> torch.Tensor:
         # Verify the generated code contains multiplication
         self.assertIn("param_0 * param_1", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_minimum(self):
         """Test associative_scan with minimum combine function."""
 
@@ -226,6 +238,9 @@ def test_min_kernel(x: torch.Tensor) -> torch.Tensor:
             "tl.minimum" in code or "triton_helpers.minimum" in code
         )
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_multiple_functions(self):
         """Test using multiple different combine functions in one kernel."""
 
@@ -262,6 +277,9 @@ def test_multi_kernel(x: torch.Tensor) -> torch.Tensor:
             "tl.maximum" in code or "triton_helpers.maximum" in code
         )
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_type_propagation(self):
         """Test that associative_scan type propagation works correctly."""
 
@@ -286,6 +304,9 @@ def test_type_kernel(x: torch.Tensor) -> torch.Tensor:
         # Use relaxed tolerance for large tensors due to accumulated floating-point errors
         torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_different_dtypes(self):
         """Test associative_scan with different data types."""
 
@@ -320,6 +341,9 @@ def test_dtype_kernel(x: torch.Tensor) -> torch.Tensor:
                         expected = expected.to(result.dtype)
                     torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_different_sizes(self):
         """Test associative_scan with different tensor sizes."""
 
@@ -356,6 +380,9 @@ def test_size_kernel(x: torch.Tensor) -> torch.Tensor:
                 expected = torch.cumsum(x, dim=1)
                 torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_reverse(self):
         """Test associative_scan with reverse=True parameter."""
 
@@ -381,6 +408,9 @@ def test_reverse_kernel(x: torch.Tensor) -> torch.Tensor:
         # Verify reverse parameter is in generated code
         self.assertIn("reverse=True", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_edge_cases(self):
         """Test associative_scan edge cases."""
 
@@ -406,6 +436,9 @@ def test_single_element(x: torch.Tensor) -> torch.Tensor:
         expected = torch.tensor([[3.0, 10.0]], device=DEVICE)
         torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_large_scale(self):
         """Test associative_scan with large tensors for performance validation."""
 
@@ -431,6 +464,9 @@ def test_large_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertEqual(result.shape, x.shape)
         self.assertEqual(result.dtype, x.dtype)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_torch_hops_mapping(self):
         """Test that torch._higher_order_ops.associative_scan automatically maps to hl.associative_scan."""
 
@@ -466,6 +502,9 @@ def test_torch_hops_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertIn("tl.associative_scan", code)
         self.assertIn("param_0 + param_1", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_code_generation(self):
         """Test that the generated code structure is correct."""
 
@@ -705,6 +744,9 @@ def cumulative_argmax_kernel(
         self.assertIn("def argmax_combine_fn_", code)
         self.assertIn("tl.associative_scan", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan is not supported by ref eager mode yet"
+    )
     def test_associative_scan_in_helper_function(self):
         """Test calling a function that internally uses hl.associative_scan."""
 
@@ -766,6 +808,7 @@ def test_cumsum_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertIn("param_0 + param_1", code)
         self.assertIn("tl.associative_scan", code)
 
+    @skipIfRefEager("hl.cumsum is not supported by ref eager mode yet")
     def test_cumsum_reverse(self):
         """Test cumsum with reverse=True."""
 
@@ -847,6 +890,7 @@ def test_cumprod_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertIn("param_0 * param_1", code)
         self.assertIn("tl.associative_scan", code)
 
+    @skipIfRefEager("hl.cumprod is not supported by ref eager mode yet")
     def test_cumprod_reverse(self):
         """Test cumprod with reverse=True."""
 
@@ -870,6 +914,7 @@ def test_cumprod_reverse_kernel(x: torch.Tensor) -> torch.Tensor:
         # Verify reverse parameter is used
         self.assertIn("reverse=True", code)
 
+    @skipIfRefEager("torch.cumprod is not supported by ref eager mode yet")
     def test_cumprod_different_dtypes(self):
         """Test cumprod with different data types."""
 
@@ -988,6 +1033,9 @@ def test_segmented_tuple_kernel(
         self.assertIn("def helion_combine_tuple_fn_", code)
         self.assertIn("tl.associative_scan", code)
 
+    @skipIfRefEager(
+        "torch._higher_order_ops.associative_scan with tuple arg is not supported by ref eager mode yet"
+    )
     def test_associative_scan_argmax_tuple_format(self):
         """Test cumulative argmax using tuple format combine function."""