jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        hyper_parameters: Namespace
        ):
        super().__init__()
        self.hp = hyper_parameters

        self.prenet = Conv1d(
            in_channels= self.hp.Encoder.Size,
            out_channels= self.hp.Decoder.Upsample.Base_Size,
            kernel_size= self.hp.Decoder.Prenet.Kernel_Size,
            padding= (self.hp.Decoder.Prenet.Kernel_Size - 1) // 2,
            # w_init_gain= 'leaky_relu'   # Don't use this line.
            )

        self.upsample_blocks = torch.nn.ModuleList()
        self.residual_blocks = torch.nn.ModuleList()
        previous_channels= self.hp.Decoder.Upsample.Base_Size
        for index, (upsample_rate, kernel_size) in enumerate(zip(
            self.hp.Decoder.Upsample.Rate,
            self.hp.Decoder.Upsample.Kernel_Size
            )):
            current_channels = self.hp.Decoder.Upsample.Base_Size // (2 ** (index + 1))
            upsample_block = torch.nn.Sequential(
                torch.nn.LeakyReLU(
                    negative_slope= self.hp.Decoder.LeakyRelu_Negative_Slope
                    ),
                torch.nn.utils.weight_norm(torch.nn.ConvTranspose1d(
                    in_channels= previous_channels,
                    out_channels= current_channels,
                    kernel_size= kernel_size,
                    stride= upsample_rate,
                    padding= (kernel_size - upsample_rate) // 2
                    ))
                )
            self.upsample_blocks.append(upsample_block)
            residual_blocks = torch.nn.ModuleList()
            for residual_kernel_size, residual_dilation_size in zip(
                self.hp.Decoder.Residual_Block.Kernel_Size,
                self.hp.Decoder.Residual_Block.Dilation_Size
                ):
                residual_blocks.append(Decoder_Residual_Block(
                    channels= current_channels,
                    kernel_size= residual_kernel_size,
                    dilations= residual_dilation_size,
                    negative_slope= self.hp.Decoder.LeakyRelu_Negative_Slope
                    ))
            self.residual_blocks.append(residual_blocks)
            previous_channels = current_channels

        self.postnet = torch.nn.Sequential(
            torch.nn.LeakyReLU(),
            torch.nn.Conv1d(
                in_channels= previous_channels,
                out_channels= 1,
                kernel_size= self.hp.Decoder.Postnet.Kernel_Size,
                padding= (self.hp.Decoder.Postnet.Kernel_Size - 1) // 2,
                bias= False
                ),
            torch.nn.Tanh(),
            Lambda(lambda x: x.squeeze(1))
            )

        # This is critical when using weight normalization.
        def weight_norm_initialize_weight(module):
            if 'Conv' in module.__class__.__name__:
                module.weight.data.normal_(0.0, 0.01)
        self.upsample_blocks.apply(weight_norm_initialize_weight)
        self.residual_blocks.apply(weight_norm_initialize_weight)
//  [python]
def __init__(
        self,
        hyper_parameters: Namespace
        ):
        super().__init__()
        self.hp = hyper_parameters

        self.prenet = Conv1d(
            in_channels= self.hp.Encoder.Size,
            out_channels= self.hp.Decoder.Upsample.Base_Size,
            kernel_size= self.hp.Decoder.Prenet.Kernel_Size,
            padding= (self.hp.Decoder.Prenet.Kernel_Size - 1) // 2,
            # w_init_gain= 'leaky_relu'   # Don't use this line.
            )

        self.upsample_blocks = torch.nn.ModuleList()
        self.residual_blocks = torch.nn.ModuleList()
        previous_channels= self.hp.Decoder.Upsample.Base_Size
        for index, (upsample_rate, kernel_size) in enumerate(zip(
            self.hp.Decoder.Upsample.Rate,
            self.hp.Decoder.Upsample.Kernel_Size
            )):
            current_channels = self.hp.Decoder.Upsample.Base_Size // (2 ** (index + 1))
            upsample_block = torch.nn.Sequential(
                torch.nn.LeakyReLU(
                    negative_slope= self.hp.Decoder.LeakyRelu_Negative_Slope
                    ),
                torch.nn.utils.parametrizations.weight_norm(torch.nn.ConvTranspose1d(
                    in_channels= previous_channels,
                    out_channels= current_channels,
                    kernel_size= kernel_size,
                    stride= upsample_rate,
                    padding= (kernel_size - upsample_rate) // 2
                    ))
                )
            self.upsample_blocks.append(upsample_block)
            residual_blocks = torch.nn.ModuleList()
            for residual_kernel_size, residual_dilation_size in zip(
                self.hp.Decoder.Residual_Block.Kernel_Size,
                self.hp.Decoder.Residual_Block.Dilation_Size
                ):
                residual_blocks.append(Decoder_Residual_Block(
                    channels= current_channels,
                    kernel_size= residual_kernel_size,
                    dilations= residual_dilation_size,
                    negative_slope= self.hp.Decoder.LeakyRelu_Negative_Slope
                    ))
            self.residual_blocks.append(residual_blocks)
            previous_channels = current_channels

        self.postnet = torch.nn.Sequential(
            torch.nn.LeakyReLU(),
            torch.nn.Conv1d(
                in_channels= previous_channels,
                out_channels= 1,
                kernel_size= self.hp.Decoder.Postnet.Kernel_Size,
                padding= (self.hp.Decoder.Postnet.Kernel_Size - 1) // 2,
                bias= False
                ),
            torch.nn.Tanh(),
            Lambda(lambda x: x.squeeze(1))
            )

        # This is critical when using weight normalization.
        def weight_norm_initialize_weight(module):
            if 'Conv' in module.__class__.__name__:
                module.weight.data.normal_(0.0, 0.01)
        self.upsample_blocks.apply(weight_norm_initialize_weight)
        self.residual_blocks.apply(weight_norm_initialize_weight)
//  [python]
def __init__(self, in_features, mask_p_t=0.1, mask_p_c=0.01, mask_t_span=6, mask_c_span=64, dropout=0.1,
                 position_encoder=25):
        super().__init__()
        self.mask_replacement = torch.nn.Parameter(torch.zeros(in_features), requires_grad=True)
        self.p_t = mask_p_t
        self.p_c = mask_p_c
        self.mask_t_span = mask_t_span
        self.mask_c_span = mask_c_span
        transformer_dim = 3 * in_features

        conv = nn.Conv1d(in_features, in_features, position_encoder, padding=position_encoder // 2, groups=16)
        nn.init.normal_(conv.weight, mean=0, std=2 / transformer_dim)
        nn.init.constant_(conv.bias, 0)
        conv = nn.utils.weight_norm(conv, dim=2)
        self.relative_position = nn.Sequential(conv, nn.GELU())

        self.input_conditioning = nn.Sequential(
            Permute([0, 2, 1]),
            nn.LayerNorm(in_features),
            nn.Dropout(dropout),
            Permute([0, 2, 1]),
            nn.Conv1d(in_features, transformer_dim, 1),
        )
//  [python]
def __init__(self, in_features, mask_p_t=0.1, mask_p_c=0.01, mask_t_span=6, mask_c_span=64, dropout=0.1,
                 position_encoder=25):
        super().__init__()
        self.mask_replacement = torch.nn.Parameter(torch.zeros(in_features), requires_grad=True)
        self.p_t = mask_p_t
        self.p_c = mask_p_c
        self.mask_t_span = mask_t_span
        self.mask_c_span = mask_c_span
        transformer_dim = 3 * in_features

        conv = nn.Conv1d(in_features, in_features, position_encoder, padding=position_encoder // 2, groups=16)
        nn.init.normal_(conv.weight, mean=0, std=2 / transformer_dim)
        nn.init.constant_(conv.bias, 0)
        conv = nn.utils.parametrizations.weight_norm(conv, dim=2)
        self.relative_position = nn.Sequential(conv, nn.GELU())

        self.input_conditioning = nn.Sequential(
            Permute([0, 2, 1]),
            nn.LayerNorm(in_features),
            nn.Dropout(dropout),
            Permute([0, 2, 1]),
            nn.Conv1d(in_features, transformer_dim, 1),
        )
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        

        C_in = args.n_colors
        n_feats = args.n_feats
        block_feats = args.block_feats
        n_layers = args.n_layers
        scale = int(args.scale)
        use_hsigmoid = args.use_hsigmoid
        use_ca = args.use_ca
        #print('scale:{0}'.format(scale))
        res_s = args.res_scale

        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
            [0.4488, 0.4371, 0.4040])).view([1, 3, 1, 1])


        wn = lambda x: torch.nn.utils.weight_norm(x)

        # define head module

        head = []
        head.append(wn(nn.Conv2d(C_in, n_feats, 3, padding=3//2, bias=True)))

        # define body module
        body = []
        self.x_scale_list = nn.ModuleList()
        self.res_scale_list = nn.ModuleList()
        self.auxilary_scale_list = nn.ModuleList()
        for _ in range(n_layers):
            body.append(Block_a(n_feats, 3, block_feats, wn=wn))
            self.x_scale_list.append(Scale(res_s))
            self.res_scale_list.append(Scale(res_s))
            self.auxilary_scale_list.append(Scale(res_s))

                    
        # define tail module
        out_feats = scale*scale*C_in
        tail = []
        tail.append(
            wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2, bias=True))
        )
        tail.append(nn.PixelShuffle(scale))

        # define skip module
        skip = []
        skip.append(
            wn(nn.Conv2d(C_in, out_feats, 3, padding=3//2, bias=True))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)

        # auxilary features
        self.fusion_conv_list = nn.ModuleList()
        for j in range(n_layers):
            if use_ca:
                tmp = nn.Sequential(*[nn.Conv2d(n_feats*(j+1), n_feats, 1, padding=0, bias=False), CALayer(n_feats, 1, use_hsigmoid=use_hsigmoid)])
            else:
                tmp = nn.Sequential(*[nn.Conv2d(n_feats*(j+1), n_feats, 1, padding=0, bias=False)])
            self.fusion_conv_list.append(tmp)

        self.test_flops = False
//  [python]
def __init__(self, args):
        super(MODEL, self).__init__()
        

        C_in = args.n_colors
        n_feats = args.n_feats
        block_feats = args.block_feats
        n_layers = args.n_layers
        scale = int(args.scale)
        use_hsigmoid = args.use_hsigmoid
        use_ca = args.use_ca
        #print('scale:{0}'.format(scale))
        res_s = args.res_scale

        self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
            [0.4488, 0.4371, 0.4040])).view([1, 3, 1, 1])


        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)

        # define head module

        head = []
        head.append(wn(nn.Conv2d(C_in, n_feats, 3, padding=3//2, bias=True)))

        # define body module
        body = []
        self.x_scale_list = nn.ModuleList()
        self.res_scale_list = nn.ModuleList()
        self.auxilary_scale_list = nn.ModuleList()
        for _ in range(n_layers):
            body.append(Block_a(n_feats, 3, block_feats, wn=wn))
            self.x_scale_list.append(Scale(res_s))
            self.res_scale_list.append(Scale(res_s))
            self.auxilary_scale_list.append(Scale(res_s))

                    
        # define tail module
        out_feats = scale*scale*C_in
        tail = []
        tail.append(
            wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2, bias=True))
        )
        tail.append(nn.PixelShuffle(scale))

        # define skip module
        skip = []
        skip.append(
            wn(nn.Conv2d(C_in, out_feats, 3, padding=3//2, bias=True))
        )
        skip.append(nn.PixelShuffle(scale))

        # make object members
        self.head = nn.Sequential(*head)
        self.body = nn.Sequential(*body)
        self.tail = nn.Sequential(*tail)
        self.skip = nn.Sequential(*skip)

        # auxilary features
        self.fusion_conv_list = nn.ModuleList()
        for j in range(n_layers):
            if use_ca:
                tmp = nn.Sequential(*[nn.Conv2d(n_feats*(j+1), n_feats, 1, padding=0, bias=False), CALayer(n_feats, 1, use_hsigmoid=use_hsigmoid)])
            else:
                tmp = nn.Sequential(*[nn.Conv2d(n_feats*(j+1), n_feats, 1, padding=0, bias=False)])
            self.fusion_conv_list.append(tmp)

        self.test_flops = False
//  [python]
def __init__(self):
        super().__init__()
        self.linear = torch.nn.utils.weight_norm(torch.nn.Linear(3, 8), dim=None)
//  [python]
def __init__(self):
        super().__init__()
        self.linear = torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(3, 8), dim=None)
//  [python]
def test_RNN_cudnn_weight_norm(self):
        input_size = 10
        hidden_size = 6
        num_layers = 2
        seq_length = 7
        batch = 6

        # runs on CPU to acquire expected output
        def check_weight_norm(m, names):
            input1 = torch.randn(seq_length, batch, input_size)
            expected_output = m(input1)

            # adds weight normalization
            m = torch.nn.utils.weight_norm(m, name=names)

            # moves to NPU
            m = m.npu()
            input1 = input1.npu()

            # otherwise, subsequent warnings will be hidden, and further tests rely on them
            warnings.simplefilter("always")
            self.assertEqual(m(input1), expected_output)

            # remove weight norm
            m = torch.nn.utils.remove_weight_norm(m, name=names)
            self.assertEqual(m(input1), expected_output)

        check_weight_norm(nn.LSTM(input_size, hidden_size, num_layers), 'weight_hh_l0')
        check_weight_norm(nn.LSTM(input_size, hidden_size, num_layers, proj_size=3), 'weight_hr_l0')
//  [python]
def test_RNN_cudnn_weight_norm(self):
        input_size = 10
        hidden_size = 6
        num_layers = 2
        seq_length = 7
        batch = 6

        # runs on CPU to acquire expected output
        def check_weight_norm(m, names):
            input1 = torch.randn(seq_length, batch, input_size)
            expected_output = m(input1)

            # adds weight normalization
            m = torch.nn.utils.parametrizations.weight_norm(m, name=names)

            # moves to NPU
            m = m.npu()
            input1 = input1.npu()

            # otherwise, subsequent warnings will be hidden, and further tests rely on them
            warnings.simplefilter("always")
            self.assertEqual(m(input1), expected_output)

            # remove weight norm
            m = torch.nn.utils.parametrizations.remove_weight_norm(m, name=names)
            self.assertEqual(m(input1), expected_output)

        check_weight_norm(nn.LSTM(input_size, hidden_size, num_layers), 'weight_hh_l0')
        check_weight_norm(nn.LSTM(input_size, hidden_size, num_layers, proj_size=3), 'weight_hr_l0')
//  [python]
def __init__(self, in_dim, out_dim, kernel_size, stride=1, padding=0, 
        output_padding=0, bias=True, weight_norm=True, scale=False, transpose=False):
        """Intializes a Conv2d augmented with weight normalization.

        (See torch.nn.utils.weight_norm for detail.)

        Args:
            in_dim: number of input channels.
            out_dim: number of output channels.
            kernel_size: size of convolving kernel.
            stride: stride of convolution.
            padding: zero-padding added to both sides of input.
            output_padding: for inferring output shape (only for transposed convolution).
            bias: True if include learnable bias parameters, False otherwise.
            weight_norm: True if apply weight normalization, False otherwise.
            scale: True if include magnitude parameters, False otherwise.
            transpose: True if transposed convolution, False otherwise.
        """
        super(WeightNormConv2d, self).__init__()
        if weight_norm:
            if transpose:
                self.conv = nn.utils.weight_norm(
                    nn.ConvTranspose2d(in_dim, out_dim, kernel_size, stride=stride, 
                        padding=padding, output_padding=output_padding, bias=bias))
            else:
                self.conv = nn.utils.weight_norm(
                    nn.Conv2d(in_dim, out_dim, kernel_size, stride=stride, 
                        padding=padding, bias=bias))
            if not scale:
                self.conv.weight_g.data = torch.ones_like(self.conv.weight_g.data)
                self.conv.weight_g.requires_grad = False    # freeze scaling
        else:
            if transpose:
                self.conv = nn.ConvTranspose2d(in_dim, out_dim, kernel_size, 
                    stride=stride, padding=padding, output_padding=output_padding, 
                    bias=bias)
            else:
                self.conv = nn.Conv2d(in_dim, out_dim, kernel_size, stride=stride, 
                    padding=padding, bias=bias)
//  [python]
def __init__(self, in_dim, out_dim, kernel_size, stride=1, padding=0, 
        output_padding=0, bias=True, weight_norm=True, scale=False, transpose=False):
        """Intializes a Conv2d augmented with weight normalization.

        (See torch.nn.utils.weight_norm for detail.)

        Args:
            in_dim: number of input channels.
            out_dim: number of output channels.
            kernel_size: size of convolving kernel.
            stride: stride of convolution.
            padding: zero-padding added to both sides of input.
            output_padding: for inferring output shape (only for transposed convolution).
            bias: True if include learnable bias parameters, False otherwise.
            weight_norm: True if apply weight normalization, False otherwise.
            scale: True if include magnitude parameters, False otherwise.
            transpose: True if transposed convolution, False otherwise.
        """
        super(WeightNormConv2d, self).__init__()
        if weight_norm:
            if transpose:
                self.conv = nn.utils.parametrizations.weight_norm(
                    nn.ConvTranspose2d(in_dim, out_dim, kernel_size, stride=stride, 
                        padding=padding, output_padding=output_padding, bias=bias))
            else:
                self.conv = nn.utils.parametrizations.weight_norm(
                    nn.Conv2d(in_dim, out_dim, kernel_size, stride=stride, 
                        padding=padding, bias=bias))
            if not scale:
                self.conv.weight_g.data = torch.ones_like(self.conv.weight_g.data)
                self.conv.weight_g.requires_grad = False    # freeze scaling
        else:
            if transpose:
                self.conv = nn.ConvTranspose2d(in_dim, out_dim, kernel_size, 
                    stride=stride, padding=padding, output_padding=output_padding, 
                    bias=bias)
            else:
                self.conv = nn.Conv2d(in_dim, out_dim, kernel_size, stride=stride, 
                    padding=padding, bias=bias)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u)//2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
        self.ups.apply(init_weights)

        if gin_channels != 0:
            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(nn.utils.parametrizations.weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u)//2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
        self.ups.apply(init_weights)

        if gin_channels != 0:
            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
//  [python]
def test_rnn_weight_norm(self):
        def check_weight_norm(lw, names, num_params):
            # This Module has 4 or 5 parameters called:
            # 'weight_ih_l0', 'weight_hh_l0', 'bias_ih_l0', 'bias_hh_l0', weight_hr_l0

            # Applying weight norm on one of them causes it to become a tensor
            lw = torch.nn.utils.weight_norm(lw, name=names)
            self.assertEqual(
                sum([isinstance(p, torch.nn.Parameter) for p in lw._flat_weights]),
                num_params - 1,
            )

            # Removing the weight norm reparametrization restores the Parameter
            lw = torch.nn.utils.remove_weight_norm(lw, name=names)
            self.assertEqual(
                sum([isinstance(p, torch.nn.Parameter) for p in lw._flat_weights]),
                num_params,
            )

            # Make sure that, upon removal of the reparametrization, the
            # `._parameters` and `.named_parameters` contain the right params.
            # Specifically, the original weight ('weight_ih_l0') should be placed
            # back in the parameters, while the reparametrization components
            # ('weight_ih_l0_v' and 'weight_ih_l0_g') should be removed.
            self.assertTrue(names in lw._parameters)
            self.assertIsNotNone(lw._parameters[names])
            self.assertTrue(names + '_v' not in lw._parameters)
            self.assertTrue(names + '_g' not in lw._parameters)
            self.assertTrue(names in dict(lw.named_parameters()))
            self.assertIsNotNone(dict(lw.named_parameters())[names])
            self.assertTrue(names + '_v' not in dict(lw.named_parameters()))
            self.assertTrue(names + '_g' not in dict(lw.named_parameters()))

        check_weight_norm(torch.nn.LSTM(32, 32), 'weight_ih_l0', 4)
        check_weight_norm(torch.nn.LSTM(32, 32, proj_size=16), 'weight_hr_l0', 5)
//  [python]
def test_rnn_weight_norm(self):
        def check_weight_norm(lw, names, num_params):
            # This Module has 4 or 5 parameters called:
            # 'weight_ih_l0', 'weight_hh_l0', 'bias_ih_l0', 'bias_hh_l0', weight_hr_l0

            # Applying weight norm on one of them causes it to become a tensor
            lw = torch.nn.utils.parametrizations.weight_norm(lw, name=names)
            self.assertEqual(
                sum([isinstance(p, torch.nn.Parameter) for p in lw._flat_weights]),
                num_params - 1,
            )

            # Removing the weight norm reparametrization restores the Parameter
            lw = torch.nn.utils.parametrizations.remove_weight_norm(lw, name=names)
            self.assertEqual(
                sum([isinstance(p, torch.nn.Parameter) for p in lw._flat_weights]),
                num_params,
            )

            # Make sure that, upon removal of the reparametrization, the
            # `._parameters` and `.named_parameters` contain the right params.
            # Specifically, the original weight ('weight_ih_l0') should be placed
            # back in the parameters, while the reparametrization components
            # ('weight_ih_l0_v' and 'weight_ih_l0_g') should be removed.
            self.assertTrue(names in lw._parameters)
            self.assertIsNotNone(lw._parameters[names])
            self.assertTrue(names + '_v' not in lw._parameters)
            self.assertTrue(names + '_g' not in lw._parameters)
            self.assertTrue(names in dict(lw.named_parameters()))
            self.assertIsNotNone(dict(lw.named_parameters())[names])
            self.assertTrue(names + '_v' not in dict(lw.named_parameters()))
            self.assertTrue(names + '_g' not in dict(lw.named_parameters()))

        check_weight_norm(torch.nn.LSTM(32, 32), 'weight_ih_l0', 4)
        check_weight_norm(torch.nn.LSTM(32, 32, proj_size=16), 'weight_hr_l0', 5)
//  [python]
def test_weight_norm(dim):

    torch.manual_seed(42)

    x = torch.randn(10)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            lin = torch.nn.Linear(10, 5)
            # Wrap the linear layer with a weight_norm - This should
            # decompose "weight" into "weight_v" and "weight_g"
            self.lin = torch.nn.utils.weight_norm(lin, "weight", dim)

        def forward(self, x):
            x = self.lin(x)
            return x, poptorch.identity_loss(x**2, reduction="sum")

    model = Model()
    weight_v_before = model.lin.weight_v.detach().clone()
    weight_g_before = model.lin.weight_g.detach().clone()

    native_out, _ = model(x)

    poptorch_model = poptorch.trainingModel(model)

    poptorch_out, _ = poptorch_model(x)

    helpers.assert_allclose(expected=native_out, actual=poptorch_out)

    tensor_names = poptorch_model.getTensorNames()
    decomposed_tensors = ["weight_v", "weight_g"]

    # Check that both decomposed tensors exist in the graph
    assert all(f"lin.{t}" in tensor_names for t in decomposed_tensors)
    # Check that they also exist in the backward graph
    assert all(f"UpdatedVar___lin.{t}" in tensor_names
               for t in decomposed_tensors)

    # Ensure that the original weight tensor does NOT exist -
    # autograd should be performed with respect to the decomposed tensors
    # only
    assert "lin.weight" not in tensor_names
    assert "UpdatedVar___lin.weight" not in tensor_names

    n = 3
    # Run a few more times to ensure that the decomposed weights are being
    # updated each time
    for i in range(n):
        weight_v_after = poptorch_model.lin.weight_v.detach().clone()
        weight_g_after = poptorch_model.lin.weight_g.detach().clone()

        # Ensure the decomposed weights changed since the previous iteration
        assert not torch.allclose(weight_v_before, weight_v_after)
        assert not torch.allclose(weight_g_before, weight_g_after)

        # Prepare for the next iteration
        if i != n - 1:
            weight_v_before = weight_v_after
            weight_g_before = weight_g_after

            poptorch_model(x)
//  [python]
def test_weight_norm(dim):

    torch.manual_seed(42)

    x = torch.randn(10)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            lin = torch.nn.Linear(10, 5)
            # Wrap the linear layer with a weight_norm - This should
            # decompose "weight" into "weight_v" and "weight_g"
            self.lin = torch.nn.utils.parametrizations.weight_norm(lin, "weight", dim)

        def forward(self, x):
            x = self.lin(x)
            return x, poptorch.identity_loss(x**2, reduction="sum")

    model = Model()
    weight_v_before = model.lin.weight_v.detach().clone()
    weight_g_before = model.lin.weight_g.detach().clone()

    native_out, _ = model(x)

    poptorch_model = poptorch.trainingModel(model)

    poptorch_out, _ = poptorch_model(x)

    helpers.assert_allclose(expected=native_out, actual=poptorch_out)

    tensor_names = poptorch_model.getTensorNames()
    decomposed_tensors = ["weight_v", "weight_g"]

    # Check that both decomposed tensors exist in the graph
    assert all(f"lin.{t}" in tensor_names for t in decomposed_tensors)
    # Check that they also exist in the backward graph
    assert all(f"UpdatedVar___lin.{t}" in tensor_names
               for t in decomposed_tensors)

    # Ensure that the original weight tensor does NOT exist -
    # autograd should be performed with respect to the decomposed tensors
    # only
    assert "lin.weight" not in tensor_names
    assert "UpdatedVar___lin.weight" not in tensor_names

    n = 3
    # Run a few more times to ensure that the decomposed weights are being
    # updated each time
    for i in range(n):
        weight_v_after = poptorch_model.lin.weight_v.detach().clone()
        weight_g_after = poptorch_model.lin.weight_g.detach().clone()

        # Ensure the decomposed weights changed since the previous iteration
        assert not torch.allclose(weight_v_before, weight_v_after)
        assert not torch.allclose(weight_g_before, weight_g_after)

        # Prepare for the next iteration
        if i != n - 1:
            weight_v_before = weight_v_after
            weight_g_before = weight_g_after

            poptorch_model(x)
//  [python]
def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int = None,
        num_layers: int = 0,
        dropout: float = 0.5,
        activation: Union[str, nn.Module, Callable] = "relu",
        batch_norm: bool = True,
        weight_norm: bool = False,
        **kwargs,
    ):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = in_dim

        if isinstance(activation, str):
            activation = get_activation(activation)
        if activation is not None and not isinstance(activation, nn.Module):
            activation = ActivationWrapper(activation)

        layers = []
        for _ in range(num_layers):
            layer = nn.Linear(in_dim, hidden_dim)
            if weight_norm:
                layer = nn.utils.weight_norm(layer, dim=-1)
            layers.append(layer)

            if batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dim))

            if activation is not None:
                layers.append(activation)

            if dropout is not None and (1.0 >= dropout >= 0.0):
                layers.append(nn.Dropout(dropout))
            in_dim = hidden_dim

        layer = nn.Linear(in_dim, out_dim)
        if weight_norm:
            layer = nn.utils.weight_norm(layer, dim=-1)
        layers.append(layer)
        self.module = nn.Sequential(*layers)
//  [python]
def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int = None,
        num_layers: int = 0,
        dropout: float = 0.5,
        activation: Union[str, nn.Module, Callable] = "relu",
        batch_norm: bool = True,
        weight_norm: bool = False,
        **kwargs,
    ):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = in_dim

        if isinstance(activation, str):
            activation = get_activation(activation)
        if activation is not None and not isinstance(activation, nn.Module):
            activation = ActivationWrapper(activation)

        layers = []
        for _ in range(num_layers):
            layer = nn.Linear(in_dim, hidden_dim)
            if weight_norm:
                layer = nn.utils.parametrizations.weight_norm(layer, dim=-1)
            layers.append(layer)

            if batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dim))

            if activation is not None:
                layers.append(activation)

            if dropout is not None and (1.0 >= dropout >= 0.0):
                layers.append(nn.Dropout(dropout))
            in_dim = hidden_dim

        layer = nn.Linear(in_dim, out_dim)
        if weight_norm:
            layer = nn.utils.parametrizations.weight_norm(layer, dim=-1)
        layers.append(layer)
        self.module = nn.Sequential(*layers)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d) or isinstance(
                    m, torch.nn.ConvTranspose1d
            ):
                torch.nn.utils.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d) or isinstance(
                    m, torch.nn.ConvTranspose1d
            ):
                torch.nn.utils.parametrizations.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.utils.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.utils.parametrizations.weight_norm(m)
                logging.debug(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(ConvTranspose1d(upsample_initial_channel // 2 ** i, upsample_initial_channel // 2 ** (i + 1), k, u, padding=(k - u) // 2)))
        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel // 2 ** (i + 1)
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))
        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
        self.ups.apply(init_weights)
        if gin_channels != 0:
            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(nn.utils.parametrizations.weight_norm(ConvTranspose1d(upsample_initial_channel // 2 ** i, upsample_initial_channel // 2 ** (i + 1), k, u, padding=(k - u) // 2)))
        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel // 2 ** (i + 1)
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))
        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
        self.ups.apply(init_weights)
        if gin_channels != 0:
            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
//  [python]
def create_model(self):
        from torch import nn
        from torch.nn.utils import weight_norm

        return weight_norm(nn.Linear(20, 40), name='weight'), None, "aten::_weight_norm"
//  [python]
def create_model(self):
        from torch import nn
        from torch.nn.utils.parametrizations import weight_norm

        return weight_norm(nn.Linear(20, 40), name='weight'), None, "aten::_weight_norm"
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        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)
            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, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        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)
            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 __init__(self, channels, kernel_size=3, dilation=(1, 3)):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
            ]
        )
        self.convs.apply(init_weights)
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[0],
                        padding=get_padding(kernel_size, dilation[0]),
                    )
                ),
                parametrizations.weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[1],
                        padding=get_padding(kernel_size, dilation[1]),
                    )
                ),
            ]
        )
        self.convs.apply(init_weights)
//  [python]
def __init__(self,
                 channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 n_layers,
                 p_dropout=0,
                 gin_channels=0,
                 mean_only=False):
        assert channels % 2 == 0, "channels should be divisible by 2"
        super().__init__()
        self.channels = channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.half_channels = channels // 2
        self.mean_only = mean_only

        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
        self.cond_layer = torch.nn.utils.weight_norm(
            torch.nn.Conv1d(
                gin_channels, 2*hidden_channels*n_layers, 1), name='weight')
        self.enc = NonCausalWaveNet(hidden_channels, kernel_size, dilation_rate,
                                    n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
        self.post = nn.Conv1d(
            hidden_channels, self.half_channels * (2 - mean_only), 1)
        self.post.weight.data.zero_()
        self.post.bias.data.zero_()
//  [python]
def __init__(self,
                 channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 n_layers,
                 p_dropout=0,
                 gin_channels=0,
                 mean_only=False):
        assert channels % 2 == 0, "channels should be divisible by 2"
        super().__init__()
        self.channels = channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.half_channels = channels // 2
        self.mean_only = mean_only

        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
        self.cond_layer = torch.nn.utils.parametrizations.weight_norm(
            torch.nn.Conv1d(
                gin_channels, 2*hidden_channels*n_layers, 1), name='weight')
        self.enc = NonCausalWaveNet(hidden_channels, kernel_size, dilation_rate,
                                    n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
        self.post = nn.Conv1d(
            hidden_channels, self.half_channels * (2 - mean_only), 1)
        self.post.weight.data.zero_()
        self.post.bias.data.zero_()
//  [python]
def __init__(self, sc_i, sc, hl, n_blocks):
        super().__init__()
        self.conv_list = nn.ModuleList()

        up_list = [hl // (sc_i * (sc ** (n_blocks-1)))] + [sc for _ in range(n_blocks - 1)]

        for u in up_list:
            convt = nn.ConvTranspose2d(1, 1, (3, 2 * u), padding=(1, u // 2), stride=(1, u))
            convt = nn.utils.weight_norm(convt)
            nn.init.kaiming_normal_(convt.weight)
            self.conv_list.append(convt)
            self.conv_list.append(nn.LeakyReLU(0.4))
//  [python]
def __init__(self, sc_i, sc, hl, n_blocks):
        super().__init__()
        self.conv_list = nn.ModuleList()

        up_list = [hl // (sc_i * (sc ** (n_blocks-1)))] + [sc for _ in range(n_blocks - 1)]

        for u in up_list:
            convt = nn.ConvTranspose2d(1, 1, (3, 2 * u), padding=(1, u // 2), stride=(1, u))
            convt = nn.utils.parametrizations.weight_norm(convt)
            nn.init.kaiming_normal_(convt.weight)
            self.conv_list.append(convt)
            self.conv_list.append(nn.LeakyReLU(0.4))
//  [python]
def __init__(
        self,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
    ):
        super().__init__()
        assert kernel_size % 2 == 1
        self.hidden_channels = hidden_channels
        self.kernel_size = (kernel_size,)
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        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)