Function _get_module

kernels/cuda/layernorm.py:313–317

()

def _get_module():
    global _module
    if _module is None:
        _module = compile_cuda(CUDA_SRC, "layernorm_cuda")
    return _module


def kernel_fn(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: