Reject deferred Tensor outputs in PMPP v2 checkers

problems/pmpp_v2/conv2d_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a 2D convolution kernel that matches the reference implementation.
   The kernel should perform 2D convolution with the given specifications
   We will benchmark different sizes, kernel sizes, channels and batch sizes but they will all 

problems/pmpp_v2/grayscale_py/task.yml

@@ -8,6 +8,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement an RGB to grayscale conversion kernel that matches the reference implementation.
   The kernel should convert square RGB images with even sizes to grayscale using the standard coefficients:
   Y = 0.2989 R + 0.5870 G + 0.1140 B

problems/pmpp_v2/histogram_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a histogram kernel that counts the number of elements falling into each bin across the specified range.
   Input values are uint8 integers in [0, 255]. All sizes are multiples of 16.
   The output must be a CUDA int64 tensor of shape (256,), containing the exact count

problems/pmpp_v2/matmul_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a custom matmul function that matches the reference implementation.
   The function should handle a tuple of input tensors and apply matmul
   The shapes of all outer and inner dimensions of tensors are multiples of 16

problems/pmpp_v2/prefixsum_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement an inclusive prefix sum (scan) kernel that matches the reference implementation.
   The kernel should compute the cumulative sum of all elements up to each position.
   Input and output use float32. Correctness is checked against float64 accumulation

problems/pmpp_v2/sort_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a sort kernel that matches the reference implementation.
   The kernel should sort the input array in ascending order using a sort algorithm of your choice.
 

problems/pmpp_v2/utils.py

@@ -28,8 +28,10 @@ def get_device(use_cuda: bool = True) -> torch.device:
 
 def _check_tensor_contract(received, expected) -> list[str]:
     """Reject malformed results before comparisons can broadcast or promote them."""
-    if not isinstance(received, torch.Tensor):
-        return ["TYPE MISMATCH: output must be a torch.Tensor"]
+    # Subclass hooks can fabricate comparisons or defer the kernel until the
+    # checker touches its result, moving computation outside the timed region.
+    if type(received) is not torch.Tensor:
+        return ["TYPE MISMATCH: output must be a plain, materialized torch.Tensor"]
     if received.shape != expected.shape:
         return [
             f"SIZE MISMATCH: got {tuple(received.shape)}, expected {tuple(expected.shape)}"

problems/pmpp_v2/vectoradd_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a float16 vector addition kernel.
 
   Input: tuple(torch.Tensor, torch.Tensor) with tensors of shape (N, N) and type torch.float16. These tensors are from

problems/pmpp_v2/vectorsum_py/task.yml

@@ -10,6 +10,9 @@ files:
 lang: "py"
 
 description: |
+  Return a plain, materialized torch.Tensor. Tensor subclasses and deferred
+  computation during checking are not allowed; compute the result inside custom_kernel.
+
   Implement a vector sum reduction kernel. This kernel computes the sum of all elements in the input tensor.
 
   Input: A tensor of shape `(N,)` with values from a normal distribution with mean 0 and variance 1.

tests/pmpp_v2/test_checkers.py

@@ -15,6 +15,24 @@ def test_reference_result_passes(problem, load_problem, examples):
     assert passed, message
 
 
+@pytest.mark.parametrize("problem", PROBLEMS)
+def test_tensor_subclasses_are_rejected_without_dispatch(
+    problem, load_problem, examples
+):
+    class InterceptedOutput(torch.Tensor):
+        @classmethod
+        def __torch_function__(cls, func, types, args=(), kwargs=None):
+            raise AssertionError("checker invoked an untrusted output hook")
+
+    ref = load_problem(problem)
+    data = examples[problem]
+    base = ref.ref_kernel(data).clone()
+    output = torch.Tensor._make_subclass(InterceptedOutput, base, False)
+    passed, message = ref.check_implementation(data, output)
+    assert not passed
+    assert "TYPE MISMATCH" in message
+
+
 @pytest.mark.parametrize("problem", PROBLEMS)
 @pytest.mark.parametrize("bad_result", ("shape", "dtype", "type", "device"))
 def test_malformed_results_are_rejected(problem, bad_result, load_problem, examples):