forked from ccf-ai-infra/TileOPs-Metax
20 lines
698 B
Python
20 lines
698 B
Python
import torch
|
|
|
|
from workloads.base import WorkloadBase
|
|
|
|
|
|
class AdaLayerNormZeroTest(WorkloadBase):
|
|
|
|
def __init__(self, m: int, n: int, dtype: torch.dtype, eps: float = 1e-5):
|
|
self.m = m
|
|
self.n = n
|
|
self.dtype = dtype
|
|
self.eps = eps
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
x = torch.randn(self.m, self.n, dtype=self.dtype, device="cuda")
|
|
scale = torch.randn(self.m, self.n, dtype=self.dtype, device="cuda")
|
|
shift = torch.randn(self.m, self.n, dtype=self.dtype, device="cuda")
|
|
gate = torch.randn(self.m, self.n, dtype=self.dtype, device="cuda")
|
|
return x, scale, shift, gate
|