Team Ai
Modelpublic

replicate/relu

sourceHugging Faceapache-2.0updated 1mo agoView on Hugging Face
0likes281downloads
benchmark.py29 linesDownload Raw Back to benchmarks
1import torch2import torch.nn.functional as F3 4from kernels.benchmark import Benchmark5 6 7class ReluBenchmark(Benchmark):8    seed: int = 429 10    def setup(self):11        self.x = torch.randn(1024, 1024, device=self.device, dtype=torch.float32)12        self.out = torch.empty_like(self.x)13 14    def benchmark_base(self):15        self.out = self.kernel.relu(self.x)16 17    def verify_base(self) -> torch.Tensor:18        return F.relu(self.x)19 20    def setup_large(self):21        self.x = torch.randn(4096, 4096, device=self.device, dtype=torch.float32)22        self.out = torch.empty_like(self.x)23 24    def benchmark_large(self):25        self.out = self.kernel.relu(self.x)26 27    def verify_large(self) -> torch.Tensor:28        return F.relu(self.x)29