[ez] update references so they run locally (#299)

* [ez] update references so they run locally

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Mark Saroufim <marksaroufim@meta.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

examples/identity_py/reference.py

@@ -1,7 +1,7 @@
 import torch
 from task import input_t, output_t
 from utils import verbose_allclose
-
+from typing import Tuple
 
 def generate_input(size: int, seed: int) -> input_t:
     gen = torch.Generator(device='cuda')
@@ -15,13 +15,13 @@ def ref_kernel(data: input_t) -> output_t:
     return data
 
 
-def check_implementation(data, output) -> str:
+def check_implementation(data: input_t, output: output_t) -> Tuple[bool, str]:
     expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
     if len(reasons) > 0:
         # TODO better processing of reasons
-        return "mismatch found! custom implementation doesn't match reference.: " + reasons[0]
+        return False, "mismatch found! custom implementation doesn't match reference: " + reasons[0]
 
-    return ''
+    return True, ''
 
 

examples/matmul_py/reference.py

@@ -1,6 +1,7 @@
 import torch
 from task import input_t, output_t
 from utils import verbose_allclose
+from typing import Tuple
 
 def generate_input(m: int, n: int, k: int, seed: int) -> input_t:
     gen = torch.Generator(device='cuda')
@@ -15,12 +16,12 @@ def ref_kernel(data: input_t) -> output_t:
     a, b = data
     return a @ b
 
-def check_implementation(data: input_t, output: output_t) -> str:
+def check_implementation(data: input_t, output: output_t) -> Tuple[bool, str]:
     expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
     if len(reasons) > 0:
         # TODO better processing of reasons
-        return "mismatch found! custom implementation doesn't match reference.: " + reasons[0]
+        return False, "mismatch found! custom implementation doesn't match reference: " + reasons[0]
 
-    return ''
+    return True, ''
 

examples/softmax_py/reference.py

@@ -1,6 +1,7 @@
 import torch
 from utils import verbose_allclose, get_device
 from task import input_t, output_t
+from typing import Tuple
 
 
 def generate_input(size: int, seed: int) -> input_t:
@@ -13,12 +14,12 @@ def generate_input(size: int, seed: int) -> input_t:
 def ref_kernel(data: input_t) -> output_t:
     return torch.nn.functional.softmax(data, dim=-1)
 
-def check_implementation(data: input_t, output: output_t) -> str:
+def check_implementation(data: input_t, output: output_t) -> Tuple[bool, str]:
 
     expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
 
     if len(reasons) > 0:
-        return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
+        return False, "mismatch found! custom implementation doesn't match reference: " + reasons[0]
 
-    return ''
+    return True, ''

examples/vectoradd_py/reference.py

@@ -1,6 +1,7 @@
 from utils import verbose_allclose
 import torch
 from task import input_t, output_t
+from typing import Tuple
 
 def ref_kernel(data: input_t) -> output_t:
     """
@@ -28,11 +29,11 @@ def generate_input(size: int, seed: int) -> input_t:
 def check_implementation(
     data: input_t,
     output: output_t,
-) -> bool:
+) -> Tuple[bool, str]:
     expected = ref_kernel(data)
     reasons = verbose_allclose(output, expected)
 
     if len(reasons) > 0:
-        return "mismatch found! custom implementation doesn't match reference: " + reasons[0]
+        return False, "mismatch found! custom implementation doesn't match reference: " + reasons[0]
 
-    return ''
+    return True, ''