nvfp4_dual_gemm / task.py
Arseni
Initialize project code, dropping binary tracking
6db16f2
Raw
History Blame Contribute Delete
369 Bytes
import torch
from typing import TypedDict, TypeVar
input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor])
output_t = TypeVar("output_t", bound=torch.Tensor)
class TestSpec(TypedDict):
m: int
n: int
k: int
l: int
seed: int