import torch.nn as nn


def zero_module(module):
    """Zero out a module's parameters and return it."""
    for param in module.parameters():
        param.detach().zero_()
    return module


def normalization(channels, channels_per_group=None):
    if channels_per_group is not None:
        return Float32GroupNorm(channels // channels_per_group, channels)

    num_groups = 4 if channels % 8 != 0 else 8
    return Float32GroupNorm(num_groups, channels)


class Float32GroupNorm(nn.GroupNorm):
    def forward(self, x):
        return super().forward(x.float()).to(dtype=x.dtype)


def conv_nd(dims, *args, **kwargs):
    if dims == 1:
        return nn.Conv1d(*args, **kwargs)
    if dims == 2:
        return nn.Conv2d(*args, **kwargs)
    if dims == 3:
        return nn.Conv3d(*args, **kwargs)
    raise ValueError(f"unsupported dimensions: {dims}")


def avg_pool_nd(dims, *args, **kwargs):
    if dims == 1:
        return nn.AvgPool1d(*args, **kwargs)
    if dims == 2:
        return nn.AvgPool2d(*args, **kwargs)
    if dims == 3:
        return nn.AvgPool3d(*args, **kwargs)
    raise ValueError(f"unsupported dimensions: {dims}")