def wang_init_(param: torch.Tensor, dim: int, num_blocks: int):
    """Adopted from https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/init_functions.py."""
    std = 2 / num_blocks / math.sqrt(dim)
    torch.nn.init.normal_(param, mean=0.0, std=std)
    return param
Adopted from https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/init_functions.py.
xlstm/components/init.py