Final

Author: S1ro1

problems/amd_distributed/gemm-rs/reference.py

@@ -1,4 +1,3 @@
-from utils import make_match_reference
 from task import input_t, output_t
 import torch
 

problems/amd_distributed/gemm-rs/task.py

@@ -1,8 +1,8 @@
-from typing import TypedDict, TypeVar, Tuple, Dict
+from typing import TypedDict, TypeVar, Tuple, Optional
 import torch
 
-input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict])
-output_t = TypeVar("output_t", bound=Tuple[torch.Tensor, Dict])
+input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]])
+output_t = TypeVar("output_t", bound=torch.Tensor)
 
 
 class TestSpec(TypedDict):
@@ -11,4 +11,4 @@ class TestSpec(TypedDict):
     n: int
     k: int
     has_bias: bool
-    seed: int
+    seed: int

problems/amd_distributed/gemm-rs/task.yml

@@ -60,6 +60,7 @@ templates:
   Python: "submission.py"
 
 ranking_by: "geom"
+ranked_timeout: 360 # just in case
 
 tests:
   - {"world_size": 8, "m": 64, "n": 2880, "k": 2880, "has_bias": True, "seed": 2035}