jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        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)
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        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)

            # 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 apply_weight_norm(self):
        self.res_block = torch.nn.utils.weight_norm(self.res_block)
        for idx, layer in enumerate(self.main_block):
            if len(layer.state_dict()) != 0:
                self.main_block[idx] = torch.nn.utils.weight_norm(layer)
        for idx, layer in enumerate(self.out_block):
            if len(layer.state_dict()) != 0:
                self.out_block[idx] = torch.nn.utils.weight_norm(layer)
//  [python]
def apply_weight_norm(self):
        self.res_block = weight_norm(self.res_block)
        for idx, layer in enumerate(self.main_block):
            if len(layer.state_dict()) != 0:
                self.main_block[idx] = weight_norm(layer)
        for idx, layer in enumerate(self.out_block):
            if len(layer.state_dict()) != 0:
                self.out_block[idx] = weight_norm(layer)
//  [python]
def __init__(self, input_size, output_size, upsample_factor, hp=None, index=0):
        super(DownsampleNet, self).__init__()
        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
        layer = nn.Conv1d(
            input_size,
            output_size,
            kernel_size=upsample_factor * 2,
            stride=upsample_factor,
            padding=upsample_factor // 2 + upsample_factor % 2,
        )

        self.layer = torch.nn.utils.weight_norm(layer)
//  [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 = torch.nn.utils.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)
//  [python]
def create_conv1d_with_weight_norm(speaker_embedding_size, hidden_size, num_layers):
    """
    Create a Conv1D layer with weight normalization.
    """
    conv1d = nn.Conv1d(speaker_embedding_size, 2 * hidden_size * num_layers, 1)
    return torch.nn.utils.weight_norm(conv1d, name='weight')
//  [python]
def create_conv1d_with_weight_norm(speaker_embedding_size, hidden_size, num_layers):
    """
    Create a Conv1D layer with weight normalization.
    """
    conv1d = nn.Conv1d(speaker_embedding_size, 2 * hidden_size * num_layers, 1)
    return parametrizations.weight_norm(conv1d, name='weight')
//  [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 = torch.nn.utils.weight_norm(self.res_conv)
        self.skip_conv = torch.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 = torch.nn.utils.weight_norm(self.filter_conv_c)
            self.gate_conv_c = torch.nn.utils.weight_norm(self.gate_conv_c)
//  [python]
def __init__(self, args):
    _logger.info("LICENSE: EDRN is implemented by IVIP-Lab. "
                 "@yyknight https://github.com/yyknight/NTIRE2019_EDRN")
    super(EDRN, self).__init__()
    r = args.scale[0]
    G0 = args.G0
    kSize = args.EDRNkSize
    rgb_mean = (0.4313, 0.4162, 0.3861)
    rgb_std = (1.0, 1.0, 1.0)
    self.sub_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std)
    wn = lambda x: torch.nn.utils.weight_norm(x)
    # number of RG blocks, conv layers, out channels
    self.D, C, G = {
      'B': (4, 10, 16),
    }[args.EDRNconfig]

    self.SFENet1 = wn(
      nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize - 1) // 2, stride=1))
    self.encoder1 = nn.Sequential(*[
      wn(nn.Conv2d(G0, 2 * G0, kSize, padding=(kSize - 1) // 2, stride=2)),
      nn.BatchNorm2d(2 * G0),
      nn.ReLU(inplace=True)])
    self.encoder2 = nn.Sequential(*[
      wn(nn.Conv2d(2 * G0, 4 * G0, kSize, padding=(kSize - 1) // 2, stride=2)),
      nn.BatchNorm2d(4 * G0),
      nn.ReLU(inplace=True)
    ])
    self.decoder1 = nn.Sequential(*[
      wn(nn.ConvTranspose2d(4 * G0, 2 * G0, 3, padding=1, output_padding=1,
                            stride=2)),
      nn.BatchNorm2d(2 * G0),
      nn.ReLU(inplace=True)])
    self.decoder2 = nn.Sequential(*[
      wn(nn.ConvTranspose2d(2 * G0, G0, 3, padding=1, output_padding=1,
                            stride=2)),
      nn.BatchNorm2d(G0),
      nn.ReLU()
    ])
    RGs0 = [
      RG(growRate0=4 * G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D)]
    RGs0.append(
      wn(nn.Conv2d(4 * G0, 4 * G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    RGs1 = [
      RG(growRate0=2 * G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D // 2)]
    RGs1.append(
      wn(nn.Conv2d(2 * G0, 2 * G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    RGs2 = [
      RG(growRate0=G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D // 4)]
    RGs2.append(
      wn(nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    self.RGs0 = nn.Sequential(*RGs0)
    self.RGs1 = nn.Sequential(*RGs1)
    self.RGs2 = nn.Sequential(*RGs2)
    self.restoration = wn(
      nn.Conv2d(G0, args.n_colors, kSize, padding=(kSize - 1) // 2, stride=1))
    self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
//  [python]
def __init__(self, args):
    _logger.info("LICENSE: EDRN is implemented by IVIP-Lab. "
                 "@yyknight https://github.com/yyknight/NTIRE2019_EDRN")
    super(EDRN, self).__init__()
    r = args.scale[0]
    G0 = args.G0
    kSize = args.EDRNkSize
    rgb_mean = (0.4313, 0.4162, 0.3861)
    rgb_std = (1.0, 1.0, 1.0)
    self.sub_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std)
    wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
    # number of RG blocks, conv layers, out channels
    self.D, C, G = {
      'B': (4, 10, 16),
    }[args.EDRNconfig]

    self.SFENet1 = wn(
      nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize - 1) // 2, stride=1))
    self.encoder1 = nn.Sequential(*[
      wn(nn.Conv2d(G0, 2 * G0, kSize, padding=(kSize - 1) // 2, stride=2)),
      nn.BatchNorm2d(2 * G0),
      nn.ReLU(inplace=True)])
    self.encoder2 = nn.Sequential(*[
      wn(nn.Conv2d(2 * G0, 4 * G0, kSize, padding=(kSize - 1) // 2, stride=2)),
      nn.BatchNorm2d(4 * G0),
      nn.ReLU(inplace=True)
    ])
    self.decoder1 = nn.Sequential(*[
      wn(nn.ConvTranspose2d(4 * G0, 2 * G0, 3, padding=1, output_padding=1,
                            stride=2)),
      nn.BatchNorm2d(2 * G0),
      nn.ReLU(inplace=True)])
    self.decoder2 = nn.Sequential(*[
      wn(nn.ConvTranspose2d(2 * G0, G0, 3, padding=1, output_padding=1,
                            stride=2)),
      nn.BatchNorm2d(G0),
      nn.ReLU()
    ])
    RGs0 = [
      RG(growRate0=4 * G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D)]
    RGs0.append(
      wn(nn.Conv2d(4 * G0, 4 * G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    RGs1 = [
      RG(growRate0=2 * G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D // 2)]
    RGs1.append(
      wn(nn.Conv2d(2 * G0, 2 * G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    RGs2 = [
      RG(growRate0=G0, nConvLayers=C, wn=wn) \
      for _ in range(self.D // 4)]
    RGs2.append(
      wn(nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1)))
    self.RGs0 = nn.Sequential(*RGs0)
    self.RGs1 = nn.Sequential(*RGs1)
    self.RGs2 = nn.Sequential(*RGs2)
    self.restoration = wn(
      nn.Conv2d(G0, args.n_colors, kSize, padding=(kSize - 1) // 2, stride=1))
    self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
//  [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 = t.nn.utils.weight_norm(layer)
        else:
            self.layer = nn.Sequential(
                nn.ReflectionPad1d(1),
                t.nn.utils.weight_norm(
                    nn.Conv1d(input_size, output_size * upsample_factor, kernel_size=3)
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                t.nn.utils.weight_norm(
                    nn.Conv1d(
                        output_size * upsample_factor,
                        output_size * upsample_factor,
                        kernel_size=3,
                    )
                ),
                nn.LeakyReLU(),
                nn.ReflectionPad1d(1),
                t.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, channels, kernel_size=3, dilation=(1, 3, 5)):
        super(ResBlock1, self).__init__()
        self.convs1 = 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]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
                ),
                weight_norm(
                    Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
                ),
                weight_norm(
                    Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
                ),
            ]
        )
        self.convs2.apply(init_weights)
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=1,
        stride=1,
        padding=None,
        dilation=1,
        bias=True,
        w_init_gain="linear",
        use_weight_norm=False,
    ):
        super(ConvNorm, self).__init__()  # pylint: disable=super-with-arguments
        if padding is None:
            assert kernel_size % 2 == 1
            padding = int(dilation * (kernel_size - 1) / 2)
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.use_weight_norm = use_weight_norm
        conv_fn = nn.Conv1d
        self.conv = conv_fn(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=bias,
        )
        nn.init.xavier_uniform_(self.conv.weight, gain=nn.init.calculate_gain(w_init_gain))
        if self.use_weight_norm:
            self.conv = torch.nn.utils.weight_norm(self.conv)
//  [python]
def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=1,
        stride=1,
        padding=None,
        dilation=1,
        bias=True,
        w_init_gain="linear",
        use_weight_norm=False,
    ):
        super(ConvNorm, self).__init__()  # pylint: disable=super-with-arguments
        if padding is None:
            assert kernel_size % 2 == 1
            padding = int(dilation * (kernel_size - 1) / 2)
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.use_weight_norm = use_weight_norm
        conv_fn = nn.Conv1d
        self.conv = conv_fn(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=bias,
        )
        nn.init.xavier_uniform_(self.conv.weight, gain=nn.init.calculate_gain(w_init_gain))
        if self.use_weight_norm:
            self.conv = nn.utils.parametrizations.weight_norm(self.conv)
//  [python]
def __init__(
        self,
        embed_dim: int,
        kernel_size: int,
        groups: int,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.kernel_size = kernel_size
        self.conv = nn.Conv1d(
            in_channels=embed_dim,
            out_channels=embed_dim,
            kernel_size=kernel_size,
            padding=kernel_size // 2,
            groups=groups,
        )

        self.conv = torch.nn.utils.weight_norm(self.conv, name="weight", dim=2)
        self.num_remove: int = 1 if kernel_size % 2 == 0 else 0
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        num_layers,
        c_in_channels=0,
        dropout_p=0,
        weight_norm=True,
    ):
        super().__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.num_layers = num_layers
        self.c_in_channels = c_in_channels
        self.dropout_p = dropout_p

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.dropout = nn.Dropout(dropout_p)

        # init conditioning layer
        if c_in_channels > 0:
            cond_layer = torch.nn.Conv1d(c_in_channels, 2 * hidden_channels * num_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
        # intermediate layers
        for i in range(num_layers):
            dilation = dilation_rate**i
            padding = int((kernel_size * dilation - dilation) / 2)
            if i == 0:
                in_layer = torch.nn.Conv1d(
                    in_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding
                )
            else:
                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 < num_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)
        # setup weight norm
        if not weight_norm:
            self.remove_weight_norm()
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        num_layers,
        c_in_channels=0,
        dropout_p=0,
        weight_norm=True,
    ):
        super().__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.num_layers = num_layers
        self.c_in_channels = c_in_channels
        self.dropout_p = dropout_p

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.dropout = nn.Dropout(dropout_p)

        # init conditioning layer
        if c_in_channels > 0:
            cond_layer = torch.nn.Conv1d(c_in_channels, 2 * hidden_channels * num_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name="weight")
        # intermediate layers
        for i in range(num_layers):
            dilation = dilation_rate**i
            padding = int((kernel_size * dilation - dilation) / 2)
            if i == 0:
                in_layer = torch.nn.Conv1d(
                    in_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding
                )
            else:
                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 < num_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)
        # setup weight norm
        if not weight_norm:
            self.remove_weight_norm()
//  [python]
def __init__(
        self,
        in_dim,
        out_dim,
        n_layers=2,
        n_channels=256,
        kernel_size=3,
        p_dropout=0.1,
        lstm_type="bilstm",
        use_linear=True,
    ):
        super(ConvLSTMLinear, self).__init__()  # pylint: disable=super-with-arguments
        self.out_dim = out_dim
        self.lstm_type = lstm_type
        self.use_linear = use_linear
        self.dropout = nn.Dropout(p=p_dropout)

        convolutions = []
        for i in range(n_layers):
            conv_layer = ConvNorm(
                in_dim if i == 0 else n_channels,
                n_channels,
                kernel_size=kernel_size,
                stride=1,
                padding=int((kernel_size - 1) / 2),
                dilation=1,
                w_init_gain="relu",
            )
            conv_layer.conv = torch.nn.utils.weight_norm(conv_layer.conv, name="weight")
            convolutions.append(conv_layer)

        self.convolutions = nn.ModuleList(convolutions)

        if not self.use_linear:
            n_channels = out_dim

        if self.lstm_type != "":
            use_bilstm = False
            lstm_channels = n_channels
            if self.lstm_type == "bilstm":
                use_bilstm = True
                lstm_channels = int(n_channels // 2)

            self.bilstm = nn.LSTM(n_channels, lstm_channels, 1, batch_first=True, bidirectional=use_bilstm)
            lstm_norm_fn_pntr = torch.nn.utils.spectral_norm
            self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0")
            if self.lstm_type == "bilstm":
                self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0_reverse")

        if self.use_linear:
            self.dense = nn.Linear(n_channels, out_dim)
//  [python]
def __init__(
        self,
        in_dim,
        out_dim,
        n_layers=2,
        n_channels=256,
        kernel_size=3,
        p_dropout=0.1,
        lstm_type="bilstm",
        use_linear=True,
    ):
        super(ConvLSTMLinear, self).__init__()  # pylint: disable=super-with-arguments
        self.out_dim = out_dim
        self.lstm_type = lstm_type
        self.use_linear = use_linear
        self.dropout = nn.Dropout(p=p_dropout)

        convolutions = []
        for i in range(n_layers):
            conv_layer = ConvNorm(
                in_dim if i == 0 else n_channels,
                n_channels,
                kernel_size=kernel_size,
                stride=1,
                padding=int((kernel_size - 1) / 2),
                dilation=1,
                w_init_gain="relu",
            )
            conv_layer = nn.utils.parametrizations.weight_norm(conv_layer.conv, name="weight")
            convolutions.append(conv_layer)

        self.convolutions = nn.ModuleList(convolutions)

        if not self.use_linear:
            n_channels = out_dim

        if self.lstm_type != "":
            use_bilstm = False
            lstm_channels = n_channels
            if self.lstm_type == "bilstm":
                use_bilstm = True
                lstm_channels = int(n_channels // 2)

            self.bilstm = nn.LSTM(n_channels, lstm_channels, 1, batch_first=True, bidirectional=use_bilstm)
            lstm_norm_fn_pntr = nn.utils.spectral_norm
            self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0")
            if self.lstm_type == "bilstm":
                self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0_reverse")

        if self.use_linear:
            self.dense = nn.Linear(n_channels, out_dim)
//  [python]
def __init__(self, kernel, channel, padding, dilations=[1, 3, 5]):
        super().__init__()
        resstack = []
        for dilation in dilations:
            resstack += [
                nn.LeakyReLU(0.2),
                nn.ReflectionPad1d(dilation),
                nn.utils.weight_norm(
                    nn.Conv1d(channel, channel, kernel_size=kernel, dilation=dilation)
                ),
                nn.LeakyReLU(0.2),
                nn.ReflectionPad1d(padding),
                nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
            ]
        self.resstack = nn.Sequential(*resstack)

        self.shortcut = nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
//  [python]
def __init__(self, kernel, channel, padding, dilations=[1, 3, 5]):
        super().__init__()
        resstack = []
        for dilation in dilations:
            resstack += [
                nn.LeakyReLU(0.2),
                nn.ReflectionPad1d(dilation),
                nn.utils.parametrizations.weight_norm(
                    nn.Conv1d(channel, channel, kernel_size=kernel, dilation=dilation)
                ),
                nn.LeakyReLU(0.2),
                nn.ReflectionPad1d(padding),
                nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
            ]
        self.resstack = nn.Sequential(*resstack)

        self.shortcut = nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
//  [python]
def __init__(  # pylint: disable=dangerous-default-value
        self,
        in_channels,
        cond_channels,
        stride,
        dilations=[1, 3, 9, 27],
        lReLU_slope=0.2,
        conv_kernel_size=3,
        cond_hop_length=256,
        kpnet_hidden_channels=64,
        kpnet_conv_size=3,
        kpnet_dropout=0.0,
    ):
        super().__init__()

        self.cond_hop_length = cond_hop_length
        self.conv_layers = len(dilations)
        self.conv_kernel_size = conv_kernel_size

        self.kernel_predictor = KernelPredictor(
            cond_channels=cond_channels,
            conv_in_channels=in_channels,
            conv_out_channels=2 * in_channels,
            conv_layers=len(dilations),
            conv_kernel_size=conv_kernel_size,
            kpnet_hidden_channels=kpnet_hidden_channels,
            kpnet_conv_size=kpnet_conv_size,
            kpnet_dropout=kpnet_dropout,
            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
        )

        self.convt_pre = nn.Sequential(
            nn.LeakyReLU(lReLU_slope),
            nn.utils.weight_norm(
                nn.ConvTranspose1d(
                    in_channels,
                    in_channels,
                    2 * stride,
                    stride=stride,
                    padding=stride // 2 + stride % 2,
                    output_padding=stride % 2,
                )
            ),
        )

        self.conv_blocks = nn.ModuleList()
        for dilation in dilations:
            self.conv_blocks.append(
                nn.Sequential(
                    nn.LeakyReLU(lReLU_slope),
                    nn.utils.weight_norm(
                        nn.Conv1d(
                            in_channels,
                            in_channels,
                            conv_kernel_size,
                            padding=dilation * (conv_kernel_size - 1) // 2,
                            dilation=dilation,
                        )
                    ),
                    nn.LeakyReLU(lReLU_slope),
                )
            )
//  [python]
def __init__(  # pylint: disable=dangerous-default-value
        self,
        in_channels,
        cond_channels,
        stride,
        dilations=[1, 3, 9, 27],
        lReLU_slope=0.2,
        conv_kernel_size=3,
        cond_hop_length=256,
        kpnet_hidden_channels=64,
        kpnet_conv_size=3,
        kpnet_dropout=0.0,
    ):
        super().__init__()

        self.cond_hop_length = cond_hop_length
        self.conv_layers = len(dilations)
        self.conv_kernel_size = conv_kernel_size

        self.kernel_predictor = KernelPredictor(
            cond_channels=cond_channels,
            conv_in_channels=in_channels,
            conv_out_channels=2 * in_channels,
            conv_layers=len(dilations),
            conv_kernel_size=conv_kernel_size,
            kpnet_hidden_channels=kpnet_hidden_channels,
            kpnet_conv_size=kpnet_conv_size,
            kpnet_dropout=kpnet_dropout,
            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
        )

        self.convt_pre = nn.Sequential(
            nn.LeakyReLU(lReLU_slope),
            nn.utils.parametrizations.weight_norm(
                nn.ConvTranspose1d(
                    in_channels,
                    in_channels,
                    2 * stride,
                    stride=stride,
                    padding=stride // 2 + stride % 2,
                    output_padding=stride % 2,
                )
            ),
        )

        self.conv_blocks = nn.ModuleList()
        for dilation in dilations:
            self.conv_blocks.append(
                nn.Sequential(
                    nn.LeakyReLU(lReLU_slope),
                    nn.utils.parametrizations.weight_norm(
                        nn.Conv1d(
                            in_channels,
                            in_channels,
                            conv_kernel_size,
                            padding=dilation * (conv_kernel_size - 1) // 2,
                            dilation=dilation,
                        )
                    ),
                    nn.LeakyReLU(lReLU_slope),
                )
            )
//  [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]),
                    )
                ),
            ]
        )
//  [python]
def test_weight_norm_pickle(self):
        m = nn.Linear(4, 5)
        m = torch.nn.utils.weight_norm(m)
        with self.assertRaisesRegex(RuntimeError, 'state_dict'):
            pickle.dumps(m)
//  [python]
def test_weight_norm_pickle(self):
        m = nn.Linear(4, 5)
        m = torch.nn.utils.parametrizations.weight_norm(m)
        with self.assertRaisesRegex(RuntimeError, 'state_dict'):
            pickle.dumps(m)
//  [python]
def test_weight_norm_deepcopy(self):
        m = nn.Linear(4, 5)
        m = torch.nn.utils.weight_norm(m)
        m2 = deepcopy(m)
        input1 = torch.randn(3, 4)
        self.assertEqual(m(input1), m2(input1))
//  [python]
def test_weight_norm_deepcopy(self):
        m = nn.Linear(4, 5)
        m = torch.nn.utils.parametrizations.weight_norm(m)
        m2 = deepcopy(m)
        input1 = torch.randn(3, 4)
        self.assertEqual(m(input1), m2(input1))
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super().__init__()
        self.convs1 = 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]),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[2],
                        padding=get_padding(kernel_size, dilation[2]),
                    )
                ),
            ]
        )

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                ),
            ]
        )
//  [python]
def __init__(
        self,
        noise_dim=64,
        channel_size=32,
        dilations=[1, 3, 9, 27],
        strides=[8, 8, 4],
        lReLU_slope=0.2,
        kpnet_conv_size=3,
        # Below are MEL configurations options that this generator requires.
        hop_length=256,
        n_mel_channels=100,
    ):
        super(UnivNetGenerator, self).__init__()
        self.mel_channel = n_mel_channels
        self.noise_dim = noise_dim
        self.hop_length = hop_length
        channel_size = channel_size
        kpnet_conv_size = kpnet_conv_size

        self.res_stack = nn.ModuleList()
        hop_length = 1
        for stride in strides:
            hop_length = stride * hop_length
            self.res_stack.append(
                LVCBlock(
                    channel_size,
                    n_mel_channels,
                    stride=stride,
                    dilations=dilations,
                    lReLU_slope=lReLU_slope,
                    cond_hop_length=hop_length,
                    kpnet_conv_size=kpnet_conv_size,
                )
            )

        self.conv_pre = nn.utils.weight_norm(
            nn.Conv1d(noise_dim, channel_size, 7, padding=3, padding_mode="reflect")
        )

        self.conv_post = nn.Sequential(
            nn.LeakyReLU(lReLU_slope),
            nn.utils.weight_norm(nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")),
            nn.Tanh(),
        )
//  [python]
def __init__(
        self,
        noise_dim=64,
        channel_size=32,
        dilations=[1, 3, 9, 27],
        strides=[8, 8, 4],
        lReLU_slope=0.2,
        kpnet_conv_size=3,
        # Below are MEL configurations options that this generator requires.
        hop_length=256,
        n_mel_channels=100,
    ):
        super(UnivNetGenerator, self).__init__()
        self.mel_channel = n_mel_channels
        self.noise_dim = noise_dim
        self.hop_length = hop_length
        channel_size = channel_size
        kpnet_conv_size = kpnet_conv_size

        self.res_stack = nn.ModuleList()
        hop_length = 1
        for stride in strides:
            hop_length = stride * hop_length
            self.res_stack.append(
                LVCBlock(
                    channel_size,
                    n_mel_channels,
                    stride=stride,
                    dilations=dilations,
                    lReLU_slope=lReLU_slope,
                    cond_hop_length=hop_length,
                    kpnet_conv_size=kpnet_conv_size,
                )
            )

        self.conv_pre = nn.utils.parametrizations.weight_norm(
            nn.Conv1d(noise_dim, channel_size, 7, padding=3, padding_mode="reflect")
        )

        self.conv_post = nn.Sequential(
            nn.LeakyReLU(lReLU_slope),
            nn.utils.parametrizations.weight_norm(nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")),
            nn.Tanh(),
        )
//  [python]
def __init__(
        self,
        in_channels,
        cond_channels,
        stride,
        dilations=[1, 3, 9, 27],
        lReLU_slope=0.2,
        conv_kernel_size=3,
        cond_hop_length=256,
        kpnet_hidden_channels=64,
        kpnet_conv_size=3,
        kpnet_dropout=0.0,
    ):
        super().__init__()

        self.cond_hop_length = cond_hop_length
        self.conv_layers = len(dilations)
        self.conv_kernel_size = conv_kernel_size

        self.kernel_predictor = KernelPredictor(
            cond_channels=cond_channels,
            conv_in_channels=in_channels,
            conv_out_channels=2 * in_channels,
            conv_layers=len(dilations),
            conv_kernel_size=conv_kernel_size,
            kpnet_hidden_channels=kpnet_hidden_channels,
            kpnet_conv_size=kpnet_conv_size,
            kpnet_dropout=kpnet_dropout,
            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
        )

        self.convt_pre = nn.Sequential(
            nn.LeakyReLU(lReLU_slope),
            weight_norm(
                nn.ConvTranspose1d(
                    in_channels,
                    in_channels,
                    2 * stride,
                    stride=stride,
                    padding=stride // 2 + stride % 2,
                    output_padding=stride % 2,
                )
            ),
        )

        self.conv_blocks = nn.ModuleList()
        for dilation in dilations:
            self.conv_blocks.append(
                nn.Sequential(
                    nn.LeakyReLU(lReLU_slope),
                    weight_norm(
                        nn.Conv1d(
                            in_channels,
                            in_channels,
                            conv_kernel_size,
                            padding=dilation * (conv_kernel_size - 1) // 2,
                            dilation=dilation,
                        )
                    ),
                    nn.LeakyReLU(lReLU_slope),
                )
            )
//  [python]
def __init__(self, nf):
        super(InvertibleConv1d, self).__init__()

        weight = torch.randn(nf, nf)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, nf):
        super(InvertibleConv1d, self).__init__()

        weight = torch.randn(nf, nf)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def test_weight_norm_parametrization(self, device):
        for dtype in [torch.float, torch.bfloat16]:
            input1 = torch.randn(3, 4, dtype=dtype, device=device)
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input1)

            # add weight normalization
            m = torch.nn.utils.weight_norm(m)
            self.assertEqual(m.weight_g.size(), m.weight.size())
            self.assertEqual(m.weight_v.size(), (5, 1))
            self.assertEqual(m(input1), expected_output)

            # remove weight norm
            torch.nn.utils.remove_weight_norm(m)
            self.assertFalse(hasattr(m, "weight_g"))
            self.assertEqual(m(input1), expected_output)

            # test with dim=1
            m = torch.nn.utils.weight_norm(m, dim=1)
            self.assertEqual(m.weight_g.size(), m.weight.size())
            self.assertEqual(m.weight_v.size(), (1, 4))
            self.assertEqual(m(input1), expected_output)

            # test with dim=None
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input1)
            m = torch.nn.utils.weight_norm(m, dim=None)
            self.assertEqual(m(input1), expected_output)
//  [python]
def test_weight_norm_parametrization(self, device):
        for dtype in [torch.float, torch.bfloat16]:
            input1 = torch.randn(3, 4, dtype=dtype, device=device)
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input1)

            # add weight normalization
            m = torch.nn.utils.parametrizations.weight_norm(m)
            self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
            self.assertEqual(m.parametrizations.weight.original0.size(), (5, 1))
            self.assertEqual(m(input1), expected_output)

            # remove weight norm
            torch.nn.utils.parametrize.remove_parametrizations(m, "weight")
            self.assertFalse(hasattr(m, "parametrizations"))
            self.assertEqual(m(input1), expected_output)

            # test with dim=1
            m = torch.nn.utils.parametrizations.weight_norm(m, dim=1)
            self.assertEqual(m.parametrizations.weight.original1.size(), m.weight.size())
            self.assertEqual(m.parametrizations.weight.original0.size(), (1, 4))
            self.assertEqual(m(input1), expected_output)

            # test with dim=None
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input1)
            m = torch.nn.utils.parametrizations.weight_norm(m, dim=None)
            self.assertEqual(m(input1), expected_output)
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
                                    bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
                                    bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def test_weight_norm_parametrization(self, device):
        for dtype in [torch.float, torch.bfloat16]:
            input = torch.randn(3, 4, dtype=dtype, device=device)
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input)

            # add weight normalization
            m = torch.nn.utils.weight_norm(m)
            self.assertEqual(
                m.weight_g.size(), m.weight.size()
            )
            self.assertEqual(m.weight_v.size(), (5, 1))
            self.assertEqual(m(input), expected_output)

            # remove weight norm
            torch.nn.utils.remove_weight_norm(m)
            self.assertFalse(hasattr(m, "weight_g"))
            self.assertEqual(m(input), expected_output)

            # test with dim=1
            m = torch.nn.utils.weight_norm(m, dim=1)
            self.assertEqual(
                m.weight_g.size(), m.weight.size()
            )
            self.assertEqual(m.weight_v.size(), (1, 4))
            self.assertEqual(m(input), expected_output)

            # test with dim=None
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input)
            m = torch.nn.utils.weight_norm(m, dim=None)
            self.assertEqual(m(input), expected_output)