jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        dilation=1,
        causal=False,
        mode="SAME",
    ):
        super(Conv, self).__init__()

        self.causal = causal
        self.mode = mode
        if self.causal and self.mode == "SAME":
            self.padding = dilation * (kernel_size - 1)
        elif self.mode == "SAME":
            self.padding = dilation * (kernel_size - 1) // 2
        else:
            self.padding = 0
        self.conv = nn.Conv1d(
            in_channels,
            out_channels,
            kernel_size,
            dilation=dilation,
            padding=self.padding,
        )
        self.conv = nn.utils.parametrizations.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)
//  [python]
def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self = torch.nn.utils.weight_norm(self)
        self.forward = super().forward
//  [python]
def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self = torch.nn.utils.parametrizations.weight_norm(self)
        self.forward = super().forward
//  [python]
def __init__(self, channel, kernel_size=3, resstack_depth=4, hp=None):
        super(ResStack, self).__init__()

        self.use_wn = Config.use_wn
        self.use_shift_scale = Config.use_shift_scale
        self.channel = channel

        def get_padding(kernel_size, dilation=1):
            return int((kernel_size * dilation - dilation) / 2)

        if self.use_shift_scale:
            self.scale_conv = nn.utils.weight_norm(
                nn.Conv1d(
                    channel, 2 * channel, kernel_size=kernel_size, dilation=1, padding=1
                )
            )

        if not self.use_wn:
            self.layers = nn.ModuleList(
                [
                    nn.Sequential(
                        nn.LeakyReLU(),
                        nn.utils.weight_norm(
                            nn.Conv1d(
                                channel,
                                channel,
                                kernel_size=kernel_size,
                                dilation=3 ** (i % 10),
                                padding=get_padding(kernel_size, 3 ** (i % 10)),
                            )
                        ),
                        nn.LeakyReLU(),
                        nn.utils.weight_norm(
                            nn.Conv1d(
                                channel,
                                channel,
                                kernel_size=kernel_size,
                                dilation=1,
                                padding=get_padding(kernel_size, 1),
                            )
                        ),
                    )
                    for i in range(resstack_depth)
                ]
            )
        else:
            self.wn = WaveNet(
                in_channels=channel,
                out_channels=channel,
                cin_channels=-1,
                num_layers=resstack_depth,
                residual_channels=channel,
                gate_channels=channel,
                skip_channels=channel,
                # kernel_size=5,
                # dilation_rate=3,
                causal=False,
                use_downup=False,
            )
//  [python]
def __init__(self, channel, kernel_size=3, resstack_depth=4, hp=None):
        super(ResStack, self).__init__()

        self.use_wn = Config.use_wn
        self.use_shift_scale = Config.use_shift_scale
        self.channel = channel

        def get_padding(kernel_size, dilation=1):
            return int((kernel_size * dilation - dilation) / 2)

        if self.use_shift_scale:
            self.scale_conv = nn.utils.parametrizations.weight_norm(
                nn.Conv1d(
                    channel, 2 * channel, kernel_size=kernel_size, dilation=1, padding=1
                )
            )

        if not self.use_wn:
            self.layers = nn.ModuleList(
                [
                    nn.Sequential(
                        nn.LeakyReLU(),
                        nn.utils.parametrizations.weight_norm(
                            nn.Conv1d(
                                channel,
                                channel,
                                kernel_size=kernel_size,
                                dilation=3 ** (i % 10),
                                padding=get_padding(kernel_size, 3 ** (i % 10)),
                            )
                        ),
                        nn.LeakyReLU(),
                        nn.utils.parametrizations.weight_norm(
                            nn.Conv1d(
                                channel,
                                channel,
                                kernel_size=kernel_size,
                                dilation=1,
                                padding=get_padding(kernel_size, 1),
                            )
                        ),
                    )
                    for i in range(resstack_depth)
                ]
            )
        else:
            self.wn = WaveNet(
                in_channels=channel,
                out_channels=channel,
                cin_channels=-1,
                num_layers=resstack_depth,
                residual_channels=channel,
                gate_channels=channel,
                skip_channels=channel,
                # kernel_size=5,
                # dilation_rate=3,
                causal=False,
                use_downup=False,
            )
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        skip_channels,
        kernel_size,
        dilation,
        cin_channels=None,
        local_conditioning=True,
        causal=False,
        mode="SAME",
    ):
        super(ResBlock, self).__init__()
        self.causal = causal
        self.local_conditioning = local_conditioning
        self.cin_channels = cin_channels
        self.mode = mode

        self.filter_conv = Conv(
            in_channels, out_channels, kernel_size, dilation, causal, mode
        )
        self.gate_conv = Conv(
            in_channels, out_channels, kernel_size, dilation, causal, mode
        )
        self.res_conv = nn.Conv1d(out_channels, in_channels, kernel_size=1)
        self.skip_conv = nn.Conv1d(out_channels, skip_channels, kernel_size=1)
        self.res_conv = nn.utils.weight_norm(self.res_conv)
        self.skip_conv = nn.utils.weight_norm(self.skip_conv)

        if self.local_conditioning:
            self.filter_conv_c = nn.Conv1d(cin_channels, out_channels, kernel_size=1)
            self.gate_conv_c = nn.Conv1d(cin_channels, out_channels, kernel_size=1)
            self.filter_conv_c = nn.utils.weight_norm(self.filter_conv_c)
            self.gate_conv_c = nn.utils.weight_norm(self.gate_conv_c)
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        skip_channels,
        kernel_size,
        dilation,
        cin_channels=None,
        local_conditioning=True,
        causal=False,
        mode="SAME",
    ):
        super(ResBlock, self).__init__()
        self.causal = causal
        self.local_conditioning = local_conditioning
        self.cin_channels = cin_channels
        self.mode = mode

        self.filter_conv = Conv(
            in_channels, out_channels, kernel_size, dilation, causal, mode
        )
        self.gate_conv = Conv(
            in_channels, out_channels, kernel_size, dilation, causal, mode
        )
        self.res_conv = nn.Conv1d(out_channels, in_channels, kernel_size=1)
        self.skip_conv = nn.Conv1d(out_channels, skip_channels, kernel_size=1)
        self.res_conv = nn.utils.parametrizations.weight_norm(self.res_conv)
        self.skip_conv = nn.utils.parametrizations.weight_norm(self.skip_conv)

        if self.local_conditioning:
            self.filter_conv_c = nn.Conv1d(cin_channels, out_channels, kernel_size=1)
            self.gate_conv_c = nn.Conv1d(cin_channels, out_channels, kernel_size=1)
            self.filter_conv_c = nn.utils.parametrizations.weight_norm(self.filter_conv_c)
            self.gate_conv_c = nn.utils.parametrizations.weight_norm(self.gate_conv_c)
//  [python]
def __init__(self, channels=16, kernel_size=3, resstack_depth=4, hp=None):
        super(ResStack2D, self).__init__()
        channels = 16
        kernel_size = 3
        resstack_depth = 2
        self.channels = channels

        def get_padding(kernel_size, dilation=1):
            return int((kernel_size * dilation - dilation) / 2)

        self.layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LeakyReLU(),
                    nn.utils.weight_norm(
                        nn.Conv2d(
                            1,
                            self.channels,
                            kernel_size,
                            dilation=(1, 3 ** (i)),
                            padding=(1, get_padding(kernel_size, 3 ** (i))),
                        )
                    ),
                    nn.LeakyReLU(),
                    nn.utils.weight_norm(
                        nn.Conv2d(
                            self.channels,
                            self.channels,
                            kernel_size,
                            dilation=(1, 3 ** (i)),
                            padding=(1, get_padding(kernel_size, 3 ** (i))),
                        )
                    ),
                    nn.LeakyReLU(),
                    nn.utils.weight_norm(nn.Conv2d(self.channels, 1, kernel_size=1)),
                )
                for i in range(resstack_depth)
            ]
        )
//  [python]
def __init__(self, channels=16, kernel_size=3, resstack_depth=4, hp=None):
        super(ResStack2D, self).__init__()
        channels = 16
        kernel_size = 3
        resstack_depth = 2
        self.channels = channels

        def get_padding(kernel_size, dilation=1):
            return int((kernel_size * dilation - dilation) / 2)

        self.layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LeakyReLU(),
                    nn.utils.parametrizations.weight_norm(
                        nn.Conv2d(
                            1,
                            self.channels,
                            kernel_size,
                            dilation=(1, 3 ** (i)),
                            padding=(1, get_padding(kernel_size, 3 ** (i))),
                        )
                    ),
                    nn.LeakyReLU(),
                    nn.utils.parametrizations.weight_norm(
                        nn.Conv2d(
                            self.channels,
                            self.channels,
                            kernel_size,
                            dilation=(1, 3 ** (i)),
                            padding=(1, get_padding(kernel_size, 3 ** (i))),
                        )
                    ),
                    nn.LeakyReLU(),
                    nn.utils.parametrizations.weight_norm(nn.Conv2d(self.channels, 1, kernel_size=1)),
                )
                for i in range(resstack_depth)
            ]
        )
//  [python]
def build(self, query_channel, key_channel, channel, nb_heads=8, dropout=0.1):
        self.to_qkv = nn.Linear(
                        query_channel + key_channel + key_channel,
                        channel*3
                    )
        self.to_qkv = torch.nn.utils.weight_norm(self.to_qkv)

        self.head_dim = channel // nb_heads
        self.nb_heads = nb_heads

        self.drop = nn.Dropout(dropout)
//  [python]
def build(self, query_channel, key_channel, channel, nb_heads=8, dropout=0.1):
        self.to_qkv = nn.Linear(
                        query_channel + key_channel + key_channel,
                        channel*3
                    )
        self.to_qkv = torch.nn.utils.parametrizations.weight_norm(self.to_qkv)

        self.head_dim = channel // nb_heads
        self.nb_heads = nb_heads

        self.drop = nn.Dropout(dropout)
//  [python]
def __init__(self, input_size, output_size, upsample_factor, hp=None, index=0):

        super(UpsampleNet, self).__init__()
        self.up_type = Config.up_type
        self.use_smooth = Config.use_smooth
        self.use_drop = Config.use_drop
        self.input_size = input_size
        self.output_size = output_size
        self.upsample_factor = upsample_factor
        self.skip_conv = nn.Conv1d(input_size, output_size, kernel_size=1)
        self.index = index
        if self.use_smooth:
            window_lens = [5, 5, 4, 3]
            self.window_len = window_lens[index]

        if self.up_type != "pn" or self.index < 3:
            # if self.up_type != "pn":
            layer = nn.ConvTranspose1d(
                input_size,
                output_size,
                upsample_factor * 2,
                upsample_factor,
                padding=upsample_factor // 2 + upsample_factor % 2,
                output_padding=upsample_factor % 2,
            )
            self.layer = nn.utils.weight_norm(layer)
        else:
            self.layer = nn.Sequential(
                nn.ReflectionPad1d(1),
                nn.utils.weight_norm(
                    nn.Conv1d(input_size, output_size * upsample_factor, kernel_size=3)
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                nn.utils.weight_norm(
                    nn.Conv1d(
                        output_size * upsample_factor,
                        output_size * upsample_factor,
                        kernel_size=3,
                    )
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                nn.utils.weight_norm(
                    nn.Conv1d(
                        output_size * upsample_factor,
                        output_size * upsample_factor,
                        kernel_size=3,
                    )
                ),
                nn.LeakyReLU(),
            )

        if hp is not None:
            self.org = Config.up_org
            self.no_skip = Config.no_skip
        else:
            self.org = False
            self.no_skip = True

        if self.use_smooth:
            self.mas = nn.Sequential(
                # LowpassBlur(output_size, self.window_len),
                MovingAverageSmooth(output_size, self.window_len),
                # MovingAverageSmooth(output_size, self.window_len),
            )
//  [python]
def __init__(self, input_size, output_size, upsample_factor, hp=None, index=0):

        super(UpsampleNet, self).__init__()
        self.up_type = Config.up_type
        self.use_smooth = Config.use_smooth
        self.use_drop = Config.use_drop
        self.input_size = input_size
        self.output_size = output_size
        self.upsample_factor = upsample_factor
        self.skip_conv = nn.Conv1d(input_size, output_size, kernel_size=1)
        self.index = index
        if self.use_smooth:
            window_lens = [5, 5, 4, 3]
            self.window_len = window_lens[index]

        if self.up_type != "pn" or self.index < 3:
            # if self.up_type != "pn":
            layer = nn.ConvTranspose1d(
                input_size,
                output_size,
                upsample_factor * 2,
                upsample_factor,
                padding=upsample_factor // 2 + upsample_factor % 2,
                output_padding=upsample_factor % 2,
            )
            self.layer = nn.utils.parametrizations.weight_norm(layer)
        else:
            self.layer = nn.Sequential(
                nn.ReflectionPad1d(1),
                nn.utils.parametrizations.weight_norm(
                    nn.Conv1d(input_size, output_size * upsample_factor, kernel_size=3)
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                nn.utils.parametrizations.weight_norm(
                    nn.Conv1d(
                        output_size * upsample_factor,
                        output_size * upsample_factor,
                        kernel_size=3,
                    )
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                nn.utils.parametrizations.weight_norm(
                    nn.Conv1d(
                        output_size * upsample_factor,
                        output_size * upsample_factor,
                        kernel_size=3,
                    )
                ),
                nn.LeakyReLU(),
            )

        if hp is not None:
            self.org = Config.up_org
            self.no_skip = Config.no_skip
        else:
            self.org = False
            self.no_skip = True

        if self.use_smooth:
            self.mas = nn.Sequential(
                # LowpassBlur(output_size, self.window_len),
                MovingAverageSmooth(output_size, self.window_len),
                # MovingAverageSmooth(output_size, self.window_len),
            )
//  [python]
def __init__(self, n_feats: int, expansion_ratio: int, res_scale: int = 1.0):
        super(ResBlock, self).__init__()
        self.res_scale = res_scale
        self.module = nn.Sequential(
            weight_norm(nn.Conv2d(n_feats + 1, n_feats * expansion_ratio, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            weight_norm(nn.Conv2d(n_feats * expansion_ratio, n_feats, kernel_size=3, padding=1))
        )
//  [python]
def __init__(self, n_feats: int, expansion_ratio: int, res_scale: int = 1.0):
        super(ResBlock, self).__init__()
        self.res_scale = res_scale
        self.module = nn.Sequential(
            parametrizations.weight_norm(nn.Conv2d(n_feats + 1, n_feats * expansion_ratio, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            parametrizations.weight_norm(nn.Conv2d(n_feats * expansion_ratio, n_feats, kernel_size=3, padding=1))
        )
//  [python]
def __init__(self, num_blocks=8, feature_dim=16, expansion=3):
        super().__init__()
        self.num_blocks = num_blocks

        # disp head
        self.in_conv = nn.Conv2d(4, feature_dim, kernel_size=3, padding=1)
        self.layers = nn.ModuleList([ResBlock(feature_dim, expansion) for _ in range(num_blocks)])
        self.out_conv = nn.Conv2d(feature_dim, 1, kernel_size=3, padding=1)

        # occ head
        self.occ_head = nn.Sequential(
            weight_norm(nn.Conv2d(1 + 3, feature_dim, kernel_size=3, padding=1)),
            weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            nn.Conv2d(feature_dim, 1, kernel_size=3, padding=1),
            nn.Sigmoid()
        )
//  [python]
def __init__(self, num_blocks=8, feature_dim=16, expansion=3):
        super().__init__()
        self.num_blocks = num_blocks

        # disp head
        self.in_conv = nn.Conv2d(4, feature_dim, kernel_size=3, padding=1)
        self.layers = nn.ModuleList([ResBlock(feature_dim, expansion) for _ in range(num_blocks)])
        self.out_conv = nn.Conv2d(feature_dim, 1, kernel_size=3, padding=1)

        # occ head
        self.occ_head = nn.Sequential(
            parametrizations.weight_norm(nn.Conv2d(1 + 3, feature_dim, kernel_size=3, padding=1)),
            parametrizations.weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            parametrizations.weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            parametrizations.weight_norm(nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1)),
            nn.ReLU(inplace=True),
            nn.Conv2d(feature_dim, 1, kernel_size=3, padding=1),
            nn.Sigmoid()
        )
//  [python]
def __init__(self, n_feats: int, expansion_ratio: int, res_scale: int = 1.0):
        super(ResBlockMotion, self).__init__()
        self.res_scale = res_scale
        self.module = nn.Sequential(
            weight_norm(nn.Conv2d(n_feats, n_feats * expansion_ratio, kernel_size=(3, 3), padding=1)),
            nn.ReLU(inplace=True),
            weight_norm(nn.Conv2d(n_feats * expansion_ratio, n_feats, kernel_size=(3, 3), padding=1))
        )
//  [python]
def __init__(self, n_feats: int, expansion_ratio: int, res_scale: int = 1.0):
        super(ResBlockMotion, self).__init__()
        self.res_scale = res_scale
        self.module = nn.Sequential(
            nn.utils.parametrizations.weight_norm(nn.Conv2d(n_feats, n_feats * expansion_ratio, kernel_size=(3, 3), padding=1)),
            nn.ReLU(inplace=True),
            nn.utils.parametrizations.weight_norm(nn.Conv2d(n_feats * expansion_ratio, n_feats, kernel_size=(3, 3), padding=1))
        )
//  [python]
def Linear(in_features, out_features, dropout=0., bias=True):
    """Weight-normalized Linear layer (input: B x T x C)"""
    m = torch.nn.Linear(in_features, out_features, bias=bias)
    m.weight.data.normal_(mean=0, std=math.sqrt((1 - dropout) / in_features))
    m.bias.data.zero_()
    return torch.nn.utils.weight_norm(m)
//  [python]
def Linear(in_features, out_features, dropout=0., bias=True):
    """Weight-normalized Linear layer (input: B x T x C)"""
    m = torch.nn.Linear(in_features, out_features, bias=bias)
    m.weight.data.normal_(mean=0, std=math.sqrt((1 - dropout) / in_features))
    m.bias.data.zero_()
    return torch.nn.utils.parametrizations.weight_norm(m)
//  [python]
def __init__(
        self,
        input_channels,
        output_channels,
        kernel_size,
        dilation):
        if penn.NORMALIZATION == 'weight':
            norm_conv = [
                torch.nn.utils.weight_norm(
                    torch.nn.Conv1d(output_channels, output_channels, 1))]
        elif penn.NORMALIZATION == 'layer':
            norm_conv = [
                torch.nn.LayerNorm(
                    (output_channels, penn.NUM_TRAINING_SAMPLES)),
                torch.nn.Conv1d(output_channels, output_channels, 1)]
        else:
            raise ValueError(
                f'Normalization method {penn.NORMALIZATION} is not defined')

        super().__init__(*(
            [
                CausalConv1d(
                    input_channels,
                    output_channels,
                    kernel_size,
                    dilation=dilation),
                torch.nn.ReLU()
            ] + norm_conv))
//  [python]
def __init__(
        self,
        input_channels,
        output_channels,
        kernel_size,
        dilation):
        if penn.NORMALIZATION == 'weight':
            norm_conv = [
                torch.nn.utils.parametrizations.weight_norm(
                    torch.nn.Conv1d(output_channels, output_channels, 1))]
        elif penn.NORMALIZATION == 'layer':
            norm_conv = [
                torch.nn.LayerNorm(
                    (output_channels, penn.NUM_TRAINING_SAMPLES)),
                torch.nn.Conv1d(output_channels, output_channels, 1)]
        else:
            raise ValueError(
                f'Normalization method {penn.NORMALIZATION} is not defined')

        super().__init__(*(
            [
                CausalConv1d(
                    input_channels,
                    output_channels,
                    kernel_size,
                    dilation=dilation),
                torch.nn.ReLU()
            ] + norm_conv))
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0

        self.hidden_channels = hidden_channels
        self.n_layers = n_layers

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding,
            )
            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0

        self.hidden_channels = hidden_channels
        self.n_layers = n_layers

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding,
            )
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def test_weight_norm(self):
        input1 = nn.Linear(20, 40).npu()
        m = torch.nn.utils.weight_norm(input1, name='weight')
        weight_g_size = m.weight_g.size()
        weight_v_size = m.weight_v.size()
        self.assertExpectedInline(str(weight_g_size), '''torch.Size([40, 1])''')
        self.assertExpectedInline(str(weight_v_size), '''torch.Size([40, 20])''')
//  [python]
def test_weight_norm(self):
        input1 = nn.Linear(20, 40).npu()
        m = torch.nn.utils.parametrizations.weight_norm(input1, name='weight')
        weight_g_size = m.weight_g.size()
        weight_v_size = m.weight_v.size()
        self.assertExpectedInline(str(weight_g_size), '''torch.Size([40, 1])''')
        self.assertExpectedInline(str(weight_v_size), '''torch.Size([40, 20])''')
//  [python]
def test_remove_weight_norm(self):
        m = torch.nn.utils.weight_norm(nn.Linear(20, 40).npu())
        output = torch.nn.utils.remove_weight_norm(m)
        self.assertExpectedInline(str(output), '''Linear(in_features=20, out_features=40, bias=True)''')
//  [python]
def test_remove_weight_norm(self):
        m = torch.nn.utils.parametrizations.weight_norm(nn.Linear(20, 40).npu())
        output = torch.nn.utils.remove_weight_norm(m)
        self.assertExpectedInline(str(output), '''Linear(in_features=20, out_features=40, bias=True)''')
//  [python]
def __init__(self, h):
        super(Generator, self).__init__()
        self.h = h
        self.num_kernels = len(h.resblock_kernel_sizes)
        self.num_upsamples = len(h.upsample_rates)
        self.m_source = SourceModuleHnNSF(sampling_rate=h.sampling_rate, harmonic_num=8)
        self.noise_convs = nn.ModuleList()
        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
        resblock = ResBlock1 if h.resblock == "1" else ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
            c_cur = h.upsample_initial_channel // (2 ** (i + 1))
            self.ups.append(weight_norm(ConvTranspose1d(h.upsample_initial_channel // (2**i), h.upsample_initial_channel // (2 ** (i + 1)), k, u, padding=(k - u) // 2)))
            if i + 1 < len(h.upsample_rates):  #
                stride_f0 = int(np.prod(h.upsample_rates[i + 1 :]))
                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
            else:
                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
        self.resblocks = nn.ModuleList()
        ch = h.upsample_initial_channel
        for i in range(len(self.ups)):
            ch //= 2
            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
                self.resblocks.append(resblock(h, ch, k, d))

        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)
        self.upp = int(np.prod(h.upsample_rates))
//  [python]
def __init__(self, h):
        super(Generator, self).__init__()
        self.h = h
        self.num_kernels = len(h.resblock_kernel_sizes)
        self.num_upsamples = len(h.upsample_rates)
        self.m_source = SourceModuleHnNSF(sampling_rate=h.sampling_rate, harmonic_num=8)
        self.noise_convs = nn.ModuleList()
        self.conv_pre = nn.utils.parametrizations.weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
        resblock = ResBlock1 if h.resblock == "1" else ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
            c_cur = h.upsample_initial_channel // (2 ** (i + 1))
            self.ups.append(nn.utils.parametrizations.weight_norm(ConvTranspose1d(h.upsample_initial_channel // (2**i), h.upsample_initial_channel // (2 ** (i + 1)), k, u, padding=(k - u) // 2)))
            if i + 1 < len(h.upsample_rates):  #
                stride_f0 = int(np.prod(h.upsample_rates[i + 1 :]))
                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
            else:
                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
        self.resblocks = nn.ModuleList()
        ch = h.upsample_initial_channel
        for i in range(len(self.ups)):
            ch //= 2
            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
                self.resblocks.append(resblock(h, ch, k, d))

        self.conv_post = nn.utils.parametrizations.weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)
        self.upp = int(np.prod(h.upsample_rates))
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        self.args = args
        scale = args.scale[0]
        n_resblocks = args.n_resblocks
        n_feats = args.n_feats
        kernel_size = 3
        act = nn.ReLU(True)
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor([args.r_mean, args.g_mean, args.b_mean])).view([1, 3, 1, 1])
        head = []
        head.append(wn(nn.Conv2d(args.n_colors, n_feats, 3, padding=3 // 2)))
        body = []
        for i in range(n_resblocks):
            body.append(Block(n_feats, kernel_size, act=act, res_scale=args.res_scale, wn=wn))
        tail = []
        out_feats = scale * scale * args.n_colors
        tail.append(wn(nn.Conv2d(n_feats, out_feats, 3, padding=3 // 2)))
        tail.append(nn.PixelShuffle(scale))
        skip = []
        skip.append(wn(nn.Conv2d(args.n_colors, out_feats, 5, padding=5 // 2)))
        skip.append(nn.PixelShuffle(scale))
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        self.args = args
        scale = args.scale[0]
        n_resblocks = args.n_resblocks
        n_feats = args.n_feats
        kernel_size = 3
        act = nn.ReLU(True)
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor([args.r_mean, args.g_mean, args.b_mean])).view([1, 3, 1, 1])
        head = []
        head.append(wn(nn.Conv2d(args.n_colors, n_feats, 3, padding=3 // 2)))
        body = []
        for i in range(n_resblocks):
            body.append(Block(n_feats, kernel_size, act=act, res_scale=args.res_scale, wn=wn))
        tail = []
        out_feats = scale * scale * args.n_colors
        tail.append(wn(nn.Conv2d(n_feats, out_feats, 3, padding=3 // 2)))
        tail.append(nn.PixelShuffle(scale))
        skip = []
        skip.append(wn(nn.Conv2d(args.n_colors, out_feats, 5, padding=5 // 2)))
        skip.append(nn.PixelShuffle(scale))
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)
//  [python]
def test_weight_norm(self):
        m = nn.LazyLinear(7)
        with self.assertRaisesRegex(ValueError, "have uninitialized parameters."):
            m = torch.nn.utils.weight_norm(m)
//  [python]
def test_weight_norm(self):
        m = nn.LazyLinear(7)
        with self.assertRaisesRegex(ValueError, "have uninitialized parameters."):
            m = torch.nn.utils.parametrizations.weight_norm(m)
//  [python]
def WeightNorm(x, dim=0):
    return nn.utils.weight_norm(x, dim=dim)
//  [python]
def WeightNorm(x, dim=0):
    return nn.utils.parametrizations.weight_norm(x, dim=dim)
//  [python]
def __init__(
        self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.weight_norm(start, name="weight")
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        for i in range(n_layers):
            dilation = 2**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                n_channels,
                2 * n_channels,
                kernel_size,
                dilation=dilation,
                padding=padding,
            )
            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels, 1)
            cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
            self.cond_layers.append(cond_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(
        self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start, name="weight")
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        for i in range(n_layers):
            dilation = 2**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                n_channels,
                2 * n_channels,
                kernel_size,
                dilation=dilation,
                padding=padding,
            )
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name="weight")
            self.in_layers.append(in_layer)

            cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels, 1)
            cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name="weight")
            self.cond_layers.append(cond_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name="weight")
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        dilation,
        batch_norm=False,
    ):
        super(CausalConvolutionBlock, self).__init__()

        # Computes left padding so that the applied convolutions are causal
        padding = (kernel_size - 1) * dilation

        layers = []

        # First causal convolution
        conv1 = torch.nn.utils.weight_norm(
            torch.nn.Conv1d(
                in_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            )
        )
        # The truncation makes the convolution causal
        chomp1 = Chomp1d(padding)
        relu1 = torch.nn.LeakyReLU()
        layers.extend([conv1, chomp1, relu1])
        if batch_norm:
            batchnorm1 = nn.BatchNorm1d(out_channels)
            layers.append(batchnorm1)

        # Second causal convolution
        conv2 = torch.nn.utils.weight_norm(
            torch.nn.Conv1d(
                out_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            )
        )
        chomp2 = Chomp1d(padding)
        relu2 = torch.nn.LeakyReLU()
        layers.extend([conv2, chomp2, relu2])
        if batch_norm:
            batchnorm2 = nn.BatchNorm1d(out_channels)
            layers.append(batchnorm2)

        # Causal network
        self.causal = torch.nn.Sequential(*layers)

        # Residual connection
        self.upordownsample = (
            torch.nn.Conv1d(in_channels, out_channels, 1)
            if in_channels != out_channels
            else None
        )
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        dilation,
        batch_norm=False,
    ):
        super(CausalConvolutionBlock, self).__init__()

        # Computes left padding so that the applied convolutions are causal
        padding = (kernel_size - 1) * dilation

        layers = []

        # First causal convolution
        conv1 = torch.nn.utils.parametrizations.weight_norm(
            torch.nn.Conv1d(
                in_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            )
        )
        # The truncation makes the convolution causal
        chomp1 = Chomp1d(padding)
        relu1 = torch.nn.LeakyReLU()
        layers.extend([conv1, chomp1, relu1])
        if batch_norm:
            batchnorm1 = nn.BatchNorm1d(out_channels)
            layers.append(batchnorm1)

        # Second causal convolution
        conv2 = torch.nn.utils.parametrizations.weight_norm(
            torch.nn.Conv1d(
                out_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            )
        )
        chomp2 = Chomp1d(padding)
        relu2 = torch.nn.LeakyReLU()
        layers.extend([conv2, chomp2, relu2])
        if batch_norm:
            batchnorm2 = nn.BatchNorm1d(out_channels)
            layers.append(batchnorm2)

        # Causal network
        self.causal = torch.nn.Sequential(*layers)

        # Residual connection
        self.upordownsample = (
            torch.nn.Conv1d(in_channels, out_channels, 1)
            if in_channels != out_channels
            else None
        )
//  [python]
def __init__(self, in_channels, out_channels, size, cond_channels, lrelu_slope=0.2):
        super().__init__()
        self.upsample = nn.UpsamplingBilinear2d(size)
        self.conv_resize = nn.utils.weight_norm(nn.Conv2d(in_channels, out_channels, 1), dim=None)
        self.conv1 = WeightNorm(
            Conv2dBias(in_channels + cond_channels, in_channels, kernel_size=3, size=size),
        )
        self.lrelu1 = nn.LeakyReLU(lrelu_slope)
        self.conv2 = WeightNorm(
            Conv2dBias(in_channels, out_channels, kernel_size=3, size=size),
        )
        self.lrelu2 = nn.LeakyReLU(lrelu_slope)
//  [python]
def __init__(self, in_channels, out_channels, size, cond_channels, lrelu_slope=0.2):
        super().__init__()
        self.upsample = nn.UpsamplingBilinear2d(size)
        self.conv_resize = nn.utils.parametrizations.weight_norm(nn.Conv2d(in_channels, out_channels, 1), dim=None)
        self.conv1 = WeightNorm(
            Conv2dBias(in_channels + cond_channels, in_channels, kernel_size=3, size=size),
        )
        self.lrelu1 = nn.LeakyReLU(lrelu_slope)
        self.conv2 = WeightNorm(
            Conv2dBias(in_channels, out_channels, kernel_size=3, size=size),
        )
        self.lrelu2 = nn.LeakyReLU(lrelu_slope)
//  [python]
def __init__(
            self,
            in_planes: int,
            out_planes: int,
            no_bias: bool,
            norm: Optional[str],
            nonlinearity: Optional[str]) -> None:

        super().__init__()
        if norm == 'wn':
            self.fc = nn.utils.weight_norm(nn.Linear(in_planes, out_planes, not no_bias))
        else:
            self.fc = nn.Linear(in_planes, out_planes, bias=not no_bias and norm is None)
        self.norm = get_norm(norm, out_planes, dim=1, affine=not no_bias)
        self.nonlinearity = get_nonlinearity(nonlinearity)
//  [python]
def __init__(
            self,
            in_planes: int,
            out_planes: int,
            no_bias: bool,
            norm: Optional[str],
            nonlinearity: Optional[str]) -> None:

        super().__init__()
        if norm == 'wn':
            self.fc = nn.utils.parametrizations.weight_norm(nn.Linear(in_planes, out_planes, not no_bias))
        else:
            self.fc = nn.Linear(in_planes, out_planes, bias=not no_bias and norm is None)
        self.norm = get_norm(norm, out_planes, dim=1, affine=not no_bias)
        self.nonlinearity = get_nonlinearity(nonlinearity)
//  [python]
def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
                 normalize=False, downsample='none', dropout_p=0.2):
        super().__init__()
        self.actv = actv
        self.normalize = normalize
        self.downsample_type = downsample
        self.learned_sc = dim_in != dim_out
        self._build_weights(dim_in, dim_out)
        self.dropout_p = dropout_p
        
        if self.downsample_type == 'none':
            self.pool = nn.Identity()
        else:
            self.pool = weight_norm(nn.Conv1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1))
//  [python]
def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
                 normalize=False, downsample='none', dropout_p=0.2):
        super().__init__()
        self.actv = actv
        self.normalize = normalize
        self.downsample_type = downsample
        self.learned_sc = dim_in != dim_out
        self._build_weights(dim_in, dim_out)
        self.dropout_p = dropout_p
        
        if self.downsample_type == 'none':
            self.pool = nn.Identity()
        else:
            self.pool = nn.utils.parametrizations.weight_norm(nn.Conv1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1))
//  [python]
def _build_weights(self, dim_in, dim_out):
        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_in, 3, 1, 1))
        self.conv2 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
        if self.normalize:
            self.norm1 = nn.InstanceNorm1d(dim_in, affine=True)
            self.norm2 = nn.InstanceNorm1d(dim_in, affine=True)
        if self.learned_sc:
            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
//  [python]
def _build_weights(self, dim_in, dim_out):
        self.conv1 = nn.utils.parametrizations.weight_norm(nn.Conv1d(dim_in, dim_in, 3, 1, 1))
        self.conv2 = nn.utils.parametrizations.weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
        if self.normalize:
            self.norm1 = nn.InstanceNorm1d(dim_in, affine=True)
            self.norm2 = nn.InstanceNorm1d(dim_in, affine=True)
        if self.learned_sc:
            self.conv1x1 = nn.utils.parametrizations.weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
//  [python]
def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
        super().__init__()
        self.embedding = nn.Embedding(n_symbols, channels)

        padding = (kernel_size - 1) // 2
        self.cnn = nn.ModuleList()
        for _ in range(depth):
            self.cnn.append(nn.Sequential(
                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
                LayerNorm(channels),
                actv,
                nn.Dropout(0.2),
            ))
        # self.cnn = nn.Sequential(*self.cnn)

        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
//  [python]
def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
        super().__init__()
        self.embedding = nn.Embedding(n_symbols, channels)

        padding = (kernel_size - 1) // 2
        self.cnn = nn.ModuleList()
        for _ in range(depth):
            self.cnn.append(nn.Sequential(
                torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
                LayerNorm(channels),
                actv,
                nn.Dropout(0.2),
            ))
        # self.cnn = nn.Sequential(*self.cnn)

        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
//  [python]
def apply_weight_norm(self):
        self.input_conv = torch.nn.utils.weight_norm(self.input_conv)
        self.output_conv = torch.nn.utils.weight_norm(self.output_conv)
//  [python]
def apply_weight_norm(self):
        self.input_conv = weight_norm(self.input_conv)
        self.output_conv = weight_norm(self.output_conv)