loosen input ranges (#46)

* principled tol

* switch dtype stuff

Author: knightron0

problems/all-pairs-shortest-path/def.py

@@ -8,6 +8,8 @@
 class all_pairs_shortest_path(Problem):
     """All-pairs shortest path problem using Floyd-Warshall algorithm."""
 
+    is_exact = True
+    
     def __init__(self):
         super().__init__(
             name="all-pairs-shortest-path"
@@ -23,7 +25,7 @@ def reference_solution(self, adj_matrix: torch.Tensor) -> torch.Tensor:
         Returns:
             Distance matrix with shortest paths between all pairs of nodes
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=adj_matrix.dtype):
             N = adj_matrix.size(0)
             device = adj_matrix.device
 

problems/argmax/def.py

@@ -7,6 +7,8 @@
 class argmax(Problem):
     """Argmax over dimension problem."""
 
+    is_exact = True
+    
     def __init__(self):
         super().__init__(
             name="argmax"
@@ -23,7 +25,7 @@ def reference_solution(self, input_tensor: torch.Tensor, dim: int) -> torch.Tens
         Returns:
             Result of argmax operation
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_tensor.dtype):
             return torch.argmax(input_tensor, dim=dim).to(torch.int32)
 
     def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:

problems/argmin/def.py

@@ -7,6 +7,8 @@
 class argmin(Problem):
     """Argmin over dimension problem."""
 
+    is_exact = True
+    
     def __init__(self):
         super().__init__(
             name="argmin"
@@ -23,7 +25,7 @@ def reference_solution(self, input_tensor: torch.Tensor, dim: int) -> torch.Tens
         Returns:
             Result of argmin operation
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_tensor.dtype):
             return torch.argmin(input_tensor, dim=dim).to(torch.int32)
 
     def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:

problems/array-sort/def.py

@@ -9,6 +9,8 @@
 class array_sort(Problem):
     """General array sorting problem (integer arrays)."""
 
+    is_exact = True
+
     def __init__(self):
         super().__init__(name="array-sort")
 

problems/avg-pool-1d/def.py

@@ -7,6 +7,8 @@
 class avg_pool_1d(Problem):
     """1D average pooling problem."""
 
+    is_exact = False
+    
     def __init__(self):
         super().__init__(
             name="avg-pool-1d"
@@ -26,7 +28,7 @@ def reference_solution(self, input_tensor: torch.Tensor, kernel_size: int,
         Returns:
             Result of average pooling
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_tensor.dtype):
             input_reshaped = input_tensor.view(1, 1, input_tensor.size(0))
 
             result = torch.nn.functional.avg_pool1d(
@@ -65,7 +67,7 @@ def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:
                 "padding": p,
                 "create_inputs": lambda size=h, kernel_size=k, stride=s, padding=p, seed=seed, dtype=dtype: (
                     *(lambda g: (
-                        torch.rand((size), device="cuda", dtype=dtype, generator=g) * 10.0 - 5.0, # uniform [-5, 5]
+                        torch.rand((size), device="cuda", dtype=dtype, generator=g) * 2.0 - 1.0, # uniform [-1, 1]
                     ))(torch.Generator(device="cuda").manual_seed(seed)),
                     kernel_size, 
                     stride, 

problems/avg-pool-2d/def.py

@@ -7,6 +7,8 @@
 class avg_pool_2d(Problem):
     """2D average pooling problem."""
 
+    is_exact = False
+    
     def __init__(self):
         super().__init__(
             name="avg-pool-2d"
@@ -26,7 +28,7 @@ def reference_solution(self, input_tensor: torch.Tensor, kernel_size: int,
         Returns:
             Result of average pooling
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_tensor.dtype):
             input_reshaped = input_tensor.view(1, 1, input_tensor.size(0), input_tensor.size(1))
 
             result = torch.nn.functional.avg_pool2d(
@@ -66,7 +68,7 @@ def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:
                 "padding": p,
                 "create_inputs": lambda h=h, w=w, k=k, s=s, p=p, seed=seed, dtype=dtype: (
                     *(lambda g: (
-                        torch.rand((h, w), device="cuda", dtype=dtype, generator=g) * 10.0 - 5.0,  # uniform [-5, 5]
+                        torch.rand((h, w), device="cuda", dtype=dtype, generator=g) * 2.0 - 1.0,  # uniform [-1, 1]
                     ))(torch.Generator(device="cuda").manual_seed(seed)),
                     k, s, p
                 )

problems/avg-pool-3d/def.py

@@ -7,6 +7,8 @@
 class avg_pool_3d(Problem):
     """3D average pooling problem."""
 
+    is_exact = False
+    
     def __init__(self):
         super().__init__(
             name="avg-pool-3d"
@@ -26,7 +28,7 @@ def reference_solution(self, input_tensor: torch.Tensor, kernel_size: int,
         Returns:
             Result of average pooling
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_tensor.dtype):
             input_reshaped = input_tensor.view(1, 1, input_tensor.size(0), input_tensor.size(1), input_tensor.size(2))
 
             result = torch.nn.functional.avg_pool3d(
@@ -66,7 +68,7 @@ def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:
                 "padding": p,
                 "create_inputs": lambda h=h, w=w, d=d, k=k, s=s, p=p, seed=seed, dtype=dtype: (
                     *(lambda g: (
-                        torch.rand((h, w, d), device="cuda", dtype=dtype, generator=g) * 10.0 - 5.0,  # uniform [-5, 5]
+                        torch.rand((h, w, d), device="cuda", dtype=dtype, generator=g) * 2.0 - 1.0,  # uniform [-1, 1]
                     ))(torch.Generator(device="cuda").manual_seed(seed)),
                     k, s, p
                 )

problems/batch-norm/def.py

@@ -8,6 +8,8 @@
 class batch_norm(Problem):
     """Batch Normalization problem."""
 
+    is_exact = False
+
     def __init__(self):
         super().__init__(
             name="batch-norm"
@@ -24,7 +26,7 @@ def reference_solution(self, x: torch.Tensor) -> torch.Tensor:
         Returns:
             torch.Tensor: Output tensor with Batch Normalization applied, same shape as input.
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=x.dtype):
             # Create BatchNorm2d layer with no affine parameters and no running stats
             bn = nn.BatchNorm2d(
                 num_features=x.size(1),  # F dimension

problems/box-blur/def.py

@@ -8,6 +8,8 @@
 class box_blur(Problem):
     """Box blur problem."""
 
+    is_exact = False
+    
     def __init__(self):
         super().__init__(
             name="box-blur"
@@ -24,7 +26,7 @@ def reference_solution(self, input_image: torch.Tensor, kernel_size: int) -> tor
         Returns:
             Blurred image of shape (height, width)
         """
-        with torch.no_grad():
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_image.dtype):
             h, w = input_image.shape
             pad = kernel_size // 2
 

problems/conv-1d/def.py

@@ -7,6 +7,8 @@
 class conv_1d(Problem):
     """1D convolution problem."""
 
+    is_exact = False
+    
     def __init__(self):
         super().__init__(
             name="conv-1d"
@@ -23,7 +25,7 @@ def reference_solution(self, input_signal: torch.Tensor, kernel: torch.Tensor) -
         Returns:
             Result of convolution with zero padding
         """
-        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=torch.float32):
+        with torch.no_grad(), torch.autocast("cuda", enabled=False, dtype=input_signal.dtype):
             # Ensure kernel size is odd
             assert kernel.size(0) % 2 == 1, "Kernel size must be odd"
 
@@ -69,7 +71,7 @@ def generate_test_cases(self, dtype: torch.dtype) -> List[Dict[str, Any]]:
                 "kernel_size": kernel_size,
                 "create_inputs": lambda s=signal_size, k=kernel_size, seed=seed, dtype=dtype: (
                     *(lambda g: (
-                        torch.rand(s, device="cuda", dtype=dtype, generator=g) * 10.0 - 5.0,
+                        torch.rand(s, device="cuda", dtype=dtype, generator=g) * 2.0 - 1.0,
                         torch.rand(k, device="cuda", dtype=dtype, generator=g) * 2.0 - 1.0,
                     ))(torch.Generator(device="cuda").manual_seed(seed)),
                 )