jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, hidden_channels, kernel_size=3, n_layers=10, dilation_rate=1,
                 p_dropout=0):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        assert (hidden_channels % 2 == 0)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.p_dropout = p_dropout

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

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

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

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

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

        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 __init__(self, frame_size, dim, q_levels, weight_norm):
        super().__init__()

        self.q_levels = q_levels

        self.embedding = torch.nn.Embedding(
            self.q_levels,
            self.q_levels
        )

        self.input = torch.nn.Conv1d(
            in_channels=q_levels,
            out_channels=dim,
            kernel_size=frame_size,
            bias=False
        )
        init.kaiming_uniform(self.input.weight)
        if weight_norm:
            self.input = torch.nn.utils.weight_norm(self.input)

        self.hidden = torch.nn.Conv1d(
            in_channels=dim,
            out_channels=dim,
            kernel_size=1
        )
        init.kaiming_uniform(self.hidden.weight)
        init.constant(self.hidden.bias, 0)
        if weight_norm:
            self.hidden = torch.nn.utils.weight_norm(self.hidden)

        self.output = torch.nn.Conv1d(
            in_channels=dim,
            out_channels=q_levels,
            kernel_size=1
        )
        nn.lecun_uniform(self.output.weight)
        init.constant(self.output.bias, 0)
        if weight_norm:
            self.output = torch.nn.utils.weight_norm(self.output)
//  [python]
def __init__(self, frame_size, dim, q_levels, weight_norm):
        super().__init__()

        self.q_levels = q_levels

        self.embedding = torch.nn.Embedding(
            self.q_levels,
            self.q_levels
        )

        self.input = torch.nn.Conv1d(
            in_channels=q_levels,
            out_channels=dim,
            kernel_size=frame_size,
            bias=False
        )
        init.kaiming_uniform(self.input.weight)
        if weight_norm:
            self.input = torch.nn.utils.parametrizations.weight_norm(self.input)

        self.hidden = torch.nn.Conv1d(
            in_channels=dim,
            out_channels=dim,
            kernel_size=1
        )
        init.kaiming_uniform(self.hidden.weight)
        init.constant(self.hidden.bias, 0)
        if weight_norm:
            self.hidden = torch.nn.utils.parametrizations.weight_norm(self.hidden)

        self.output = torch.nn.Conv1d(
            in_channels=dim,
            out_channels=q_levels,
            kernel_size=1
        )
        nn.lecun_uniform(self.output.weight)
        init.constant(self.output.bias, 0)
        if weight_norm:
            self.output = torch.nn.utils.parametrizations.weight_norm(self.output)
//  [python]
def weight_norm_(module, class_or_tuple=nn.Conv2d, recursive=True, name='weight', add_repr=False, **kwargs):
    """Weight normalization.

    Applies weight normalization to parameters of all occurrences of ``class_or_tuple`` in the given module.

    Note:
        This is an inplace operation.

    References:
        - https://proceedings.neurips.cc/paper/2016/file/ed265bc903a5a097f61d3ec064d96d2e-Paper.pdf

    Args:
        module: Module.
        class_or_tuple: Class or tuple of classes whose parameters are to be normalized.
        recursive: Whether to search for modules recursively.
        name: Name of weight parameter.
        add_repr: Whether to indicate use of weight norm in a module's representation.
            Note that this may impair pickling.
        **kwargs: Additional keyword arguments for ``torch.nn.utils.weight_norm``.
    """

    def extra_repr(self):
        if 'torch.nn.utils.weight_norm.WeightNorm' in str(list(self._forward_pre_hooks.values())):
            return 'weight_norm=True'

    for handle, k, mod in iter_submodules(module, class_or_tuple, recursive=recursive):
        if mod._parameters.get(name) is not None:
            handle[k] = nn.utils.weight_norm(handle[k], name=name, **kwargs)
            if add_repr:
                inject_extra_repr_(handle[k], 'weight_norm', extra_repr)
//  [python]
def weight_norm_(module, class_or_tuple=nn.Conv2d, recursive=True, name='weight', add_repr=False, **kwargs):
    """Weight normalization.

    Applies weight normalization to parameters of all occurrences of ``class_or_tuple`` in the given module.

    Note:
        This is an inplace operation.

    References:
        - https://proceedings.neurips.cc/paper/2016/file/ed265bc903a5a097f61d3ec064d96d2e-Paper.pdf

    Args:
        module: Module.
        class_or_tuple: Class or tuple of classes whose parameters are to be normalized.
        recursive: Whether to search for modules recursively.
        name: Name of weight parameter.
        add_repr: Whether to indicate use of weight norm in a module's representation.
            Note that this may impair pickling.
        **kwargs: Additional keyword arguments for ``torch.nn.utils.weight_norm``.
    """

    def extra_repr(self):
        if 'torch.nn.utils.parametrizations.weight_norm.WeightNorm' in str(list(self._forward_pre_hooks.values())):
            return 'weight_norm=True'

    for handle, k, mod in iter_submodules(module, class_or_tuple, recursive=recursive):
        if mod._parameters.get(name) is not None:
            handle[k] = nn.utils.parametrizations.weight_norm(handle[k], name=name, **kwargs)
            if add_repr:
                inject_extra_repr_(handle[k], 'weight_norm', extra_repr)
//  [python]
def _create_layers(self, mlp=False):
        self.item_class_em = nn.Embedding(337, 8, max_norm=8, norm_type=2)
        self.item_family_em = nn.Embedding(33, 8, max_norm=8, norm_type=2)
        self.store_type_em = nn.Embedding(5, 2, max_norm=2, norm_type=2)
        self.store_cluster_em = nn.Embedding(17, 5, max_norm=5, norm_type=2)
        self.store_city_em = nn.Embedding(22, 5, max_norm=5, norm_type=2)
        self.store_em = nn.Embedding(54, 5, max_norm=5, norm_type=2)
        self.weekday_em = nn.Embedding(7, 3, max_norm=3, norm_type=2)
        self.day_em = nn.Embedding(31, 10, max_norm=10, norm_type=2)
        self.month_em = nn.Embedding(6, 2, max_norm=2, norm_type=2)
        self.year_em = nn.Embedding(3, 2, max_norm=2, norm_type=2)
        if not mlp:
            self.step_one_network = nn.Sequential(
                torch.nn.utils.weight_norm(nn.Linear(
                    self.hidden_size, self.hidden_size)),
                nn.ReLU(),
                nn.Dropout(self.odrop),
                torch.nn.utils.weight_norm(nn.Linear(self.hidden_size, 1))
            )
//  [python]
def _create_layers(self, mlp=False):
        self.item_class_em = nn.Embedding(337, 8, max_norm=8, norm_type=2)
        self.item_family_em = nn.Embedding(33, 8, max_norm=8, norm_type=2)
        self.store_type_em = nn.Embedding(5, 2, max_norm=2, norm_type=2)
        self.store_cluster_em = nn.Embedding(17, 5, max_norm=5, norm_type=2)
        self.store_city_em = nn.Embedding(22, 5, max_norm=5, norm_type=2)
        self.store_em = nn.Embedding(54, 5, max_norm=5, norm_type=2)
        self.weekday_em = nn.Embedding(7, 3, max_norm=3, norm_type=2)
        self.day_em = nn.Embedding(31, 10, max_norm=10, norm_type=2)
        self.month_em = nn.Embedding(6, 2, max_norm=2, norm_type=2)
        self.year_em = nn.Embedding(3, 2, max_norm=2, norm_type=2)
        if not mlp:
            self.step_one_network = nn.Sequential(
                torch.nn.utils.parametrizations.weight_norm(nn.Linear(
                    self.hidden_size, self.hidden_size)),
                nn.ReLU(),
                nn.Dropout(self.odrop),
                torch.nn.utils.parametrizations.weight_norm(nn.Linear(self.hidden_size, 1))
            )
//  [python]
def __init__(self, frame_size, n_frame_samples, n_rnn, dim,
                 learn_h0, weight_norm):
        super().__init__()

        self.frame_size = frame_size
        self.n_frame_samples = n_frame_samples
        self.dim = dim

        h0 = torch.zeros(n_rnn, dim)
        if learn_h0:
            self.h0 = torch.nn.Parameter(h0)
        else:
            self.register_buffer('h0', torch.autograd.Variable(h0))

        self.input_expand = torch.nn.Conv1d(
            in_channels=n_frame_samples,
            out_channels=dim,
            kernel_size=1
        )
        init.kaiming_uniform(self.input_expand.weight)
        init.constant(self.input_expand.bias, 0)
        if weight_norm:
            self.input_expand = torch.nn.utils.weight_norm(self.input_expand)

        self.rnn = torch.nn.GRU(
            input_size=dim,
            hidden_size=dim,
            num_layers=n_rnn,
            batch_first=True
        )
        for i in range(n_rnn):
            nn.concat_init(
                getattr(self.rnn, 'weight_ih_l{}'.format(i)),
                [nn.lecun_uniform, nn.lecun_uniform, nn.lecun_uniform]
            )
            init.constant(getattr(self.rnn, 'bias_ih_l{}'.format(i)), 0)

            nn.concat_init(
                getattr(self.rnn, 'weight_hh_l{}'.format(i)),
                [nn.lecun_uniform, nn.lecun_uniform, init.orthogonal]
            )
            init.constant(getattr(self.rnn, 'bias_hh_l{}'.format(i)), 0)

        self.upsampling = nn.LearnedUpsampling1d(
            in_channels=dim,
            out_channels=dim,
            kernel_size=frame_size
        )
        init.uniform(
            self.upsampling.conv_t.weight, -np.sqrt(6 / dim), np.sqrt(6 / dim)
        )
        init.constant(self.upsampling.bias, 0)
        if weight_norm:
            self.upsampling.conv_t = torch.nn.utils.weight_norm(
                self.upsampling.conv_t
            )
//  [python]
def __init__(self, frame_size, n_frame_samples, n_rnn, dim,
                 learn_h0, weight_norm):
        super().__init__()

        self.frame_size = frame_size
        self.n_frame_samples = n_frame_samples
        self.dim = dim

        h0 = torch.zeros(n_rnn, dim)
        if learn_h0:
            self.h0 = torch.nn.Parameter(h0)
        else:
            self.register_buffer('h0', torch.autograd.Variable(h0))

        self.input_expand = torch.nn.Conv1d(
            in_channels=n_frame_samples,
            out_channels=dim,
            kernel_size=1
        )
        init.kaiming_uniform(self.input_expand.weight)
        init.constant(self.input_expand.bias, 0)
        if weight_norm:
            self.input_expand = torch.nn.utils.parametrizations.weight_norm(self.input_expand)

        self.rnn = torch.nn.GRU(
            input_size=dim,
            hidden_size=dim,
            num_layers=n_rnn,
            batch_first=True
        )
        for i in range(n_rnn):
            nn.concat_init(
                getattr(self.rnn, 'weight_ih_l{}'.format(i)),
                [nn.lecun_uniform, nn.lecun_uniform, nn.lecun_uniform]
            )
            init.constant(getattr(self.rnn, 'bias_ih_l{}'.format(i)), 0)

            nn.concat_init(
                getattr(self.rnn, 'weight_hh_l{}'.format(i)),
                [nn.lecun_uniform, nn.lecun_uniform, init.orthogonal]
            )
            init.constant(getattr(self.rnn, 'bias_hh_l{}'.format(i)), 0)

        self.upsampling = nn.LearnedUpsampling1d(
            in_channels=dim,
            out_channels=dim,
            kernel_size=frame_size
        )
        init.uniform(
            self.upsampling.conv_t.weight, -np.sqrt(6 / dim), np.sqrt(6 / dim)
        )
        init.constant(self.upsampling.bias, 0)
        if weight_norm:
            self.upsampling.conv_t = torch.nn.utils.parametrizations.weight_norm(
                self.upsampling.conv_t
            )
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size, device, logger):
        super(WN, self).__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.device = device
        self.logger = logger
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

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

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

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

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

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1).to(self.device)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight').to(self.device)
            self.res_skip_layers.append(res_skip_layer)

        self.in_layers = self.in_layers.to(self.device)
        self.cond_layers = self.cond_layers.to(self.device)
        self.res_skip_layers = self.res_skip_layers.to(self.device)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size, device, logger):
        super(WN, self).__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.device = device
        self.logger = logger
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.cond_layers = torch.nn.ModuleList()

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

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

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

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

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1).to(self.device)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight').to(self.device)
            self.res_skip_layers.append(res_skip_layer)

        self.in_layers = self.in_layers.to(self.device)
        self.cond_layers = self.cond_layers.to(self.device)
        self.res_skip_layers = self.res_skip_layers.to(self.device)
//  [python]
def __init__(
            self, res_channels, skip_channels, dilation=1,
            diffusion_step_embed_dim_out=512,
            unconditional=True,
            mel_upsample=[16,16],
        ):
        super(Residual_block, self).__init__()
        self.res_channels = res_channels

        # the layer-specific fc for diffusion step embedding
        self.fc_t = nn.Linear(diffusion_step_embed_dim_out, self.res_channels)

        # dilated conv layer
        self.dilated_conv_layer = Conv(self.res_channels, 2 * self.res_channels, kernel_size=3, dilation=dilation)

        self.unconditional = unconditional
        if not self.unconditional:
            # add mel spectrogram upsampler and conditioner conv1x1 layer
            self.upsample_conv2d = torch.nn.ModuleList()
            for s in mel_upsample:
                conv_trans2d = torch.nn.ConvTranspose2d(1, 1, (3, 2 * s), padding=(1, s // 2), stride=(1, s))
                conv_trans2d = torch.nn.utils.weight_norm(conv_trans2d)
                torch.nn.init.kaiming_normal_(conv_trans2d.weight)
                self.upsample_conv2d.append(conv_trans2d)
            self.mel_conv = Conv(80, 2 * self.res_channels, kernel_size=1)  # 80 is mel bands

        # residual conv1x1 layer, connect to next residual layer
        self.res_conv = nn.Conv1d(res_channels, res_channels, kernel_size=1)
        self.res_conv = nn.utils.weight_norm(self.res_conv)
        nn.init.kaiming_normal_(self.res_conv.weight)

        # skip conv1x1 layer, add to all skip outputs through skip connections
        self.skip_conv = nn.Conv1d(res_channels, skip_channels, kernel_size=1)
        self.skip_conv = nn.utils.weight_norm(self.skip_conv)
        nn.init.kaiming_normal_(self.skip_conv.weight)
//  [python]
def __init__(
            self, res_channels, skip_channels, dilation=1,
            diffusion_step_embed_dim_out=512,
            unconditional=True,
            mel_upsample=[16,16],
        ):
        super(Residual_block, self).__init__()
        self.res_channels = res_channels

        # the layer-specific fc for diffusion step embedding
        self.fc_t = nn.Linear(diffusion_step_embed_dim_out, self.res_channels)

        # dilated conv layer
        self.dilated_conv_layer = Conv(self.res_channels, 2 * self.res_channels, kernel_size=3, dilation=dilation)

        self.unconditional = unconditional
        if not self.unconditional:
            # add mel spectrogram upsampler and conditioner conv1x1 layer
            self.upsample_conv2d = torch.nn.ModuleList()
            for s in mel_upsample:
                conv_trans2d = torch.nn.ConvTranspose2d(1, 1, (3, 2 * s), padding=(1, s // 2), stride=(1, s))
                conv_trans2d = torch.nn.utils.parametrizations.weight_norm(conv_trans2d)
                torch.nn.init.kaiming_normal_(conv_trans2d.weight)
                self.upsample_conv2d.append(conv_trans2d)
            self.mel_conv = Conv(80, 2 * self.res_channels, kernel_size=1)  # 80 is mel bands

        # residual conv1x1 layer, connect to next residual layer
        self.res_conv = nn.Conv1d(res_channels, res_channels, kernel_size=1)
        self.res_conv = torch.nn.utils.parametrizations.weight_norm(self.res_conv)
        nn.init.kaiming_normal_(self.res_conv.weight)

        # skip conv1x1 layer, add to all skip outputs through skip connections
        self.skip_conv = nn.Conv1d(res_channels, skip_channels, kernel_size=1)
        self.skip_conv = torch.nn.utils.parametrizations.weight_norm(self.skip_conv)
        nn.init.kaiming_normal_(self.skip_conv.weight)
//  [python]
def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, wnorm=False):
        super().__init__(in_features=in_features,
                         out_features=out_features,
                         bias=bias,
                         device=device,
                         dtype=dtype)
        if wnorm:
            weight_norm(self)

        self._fix_weight_norm_deepcopy()
//  [python]
def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, wnorm=False):
        super().__init__(in_features=in_features,
                         out_features=out_features,
                         bias=bias,
                         device=device,
                         dtype=dtype)
        if wnorm:
            torch.nn.utils.parametrizations.weight_norm(self)

        self._fix_weight_norm_deepcopy()
//  [python]
def __init__(
            self, cnn_hidden_size, rnn_hidden_size, skip_hidden_size,
            hdrop, edrop, odrop, skip=7, cnn_kernel=3, *,
            y_scale_by, steps=15, rnn_type="GRU", min_length=TRAIN_PERIODS,
            ar_window_size=0):
        super(LSTNet, self).__init__(
            bidirectional=False, edrop=edrop, odrop=odrop, propagate=False,
            min_length=min_length, y_scale_by=y_scale_by, steps=steps
        )
        self.hdrop = hdrop
        self.hdropout = LockedDropout(batch_first=False)
        self.ar_window_size = ar_window_size
        self.skip = skip
        self.skip_seq_len = (TRAIN_PERIODS - cnn_kernel) // self.skip
        self._create_layers(mlp=True)
        self.cnn = nn.utils.weight_norm(
            nn.Conv1d(self.input_dim, cnn_hidden_size, cnn_kernel))
        if rnn_type == "GRU":
            self.rnn = nn.GRU(
                cnn_hidden_size, rnn_hidden_size, batch_first=False)
            self.rnn = weight_norm_rnn(self.rnn)
        else:
            self.rnn = nn.LSTM(
                cnn_hidden_size, rnn_hidden_size, batch_first=False)
        if skip:
            if rnn_type == "GRU":
                self.skiprnn = nn.GRU(
                    cnn_hidden_size, skip_hidden_size, batch_first=False)
                self.skiprnn = weight_norm_rnn(self.skiprnn)
            elif rnn_type == "SRU":
                self.skiprnn = nn.SRU(cnn_hidden_size, skip_hidden_size)
            else:
                self.skiprnn = nn.LSTM(
                    cnn_hidden_size, skip_hidden_size, batch_first=False)
            output_dim = rnn_hidden_size + \
                skip_hidden_size * skip + steps * 7
        else:
            output_dim = rnn_hidden_size + steps * 7
        self.output_network = nn.Sequential(
            torch.nn.utils.weight_norm(
                nn.Linear(output_dim, rnn_hidden_size)),
            nn.ReLU(),
            nn.Dropout(self.odrop),
            torch.nn.utils.weight_norm(
                nn.Linear(rnn_hidden_size, steps + 1))
        )
        if self.ar_window_size:
            self.ar_highway = torch.nn.utils.weight_norm(
                nn.Linear(ar_window_size, steps + 1))
        self.init_mlp_weights()
//  [python]
def __init__(
            self, cnn_hidden_size, rnn_hidden_size, skip_hidden_size,
            hdrop, edrop, odrop, skip=7, cnn_kernel=3, *,
            y_scale_by, steps=15, rnn_type="GRU", min_length=TRAIN_PERIODS,
            ar_window_size=0):
        super(LSTNet, self).__init__(
            bidirectional=False, edrop=edrop, odrop=odrop, propagate=False,
            min_length=min_length, y_scale_by=y_scale_by, steps=steps
        )
        self.hdrop = hdrop
        self.hdropout = LockedDropout(batch_first=False)
        self.ar_window_size = ar_window_size
        self.skip = skip
        self.skip_seq_len = (TRAIN_PERIODS - cnn_kernel) // self.skip
        self._create_layers(mlp=True)
        self.cnn = nn.utils.parametrizations.weight_norm(
            nn.Conv1d(self.input_dim, cnn_hidden_size, cnn_kernel))
        if rnn_type == "GRU":
            self.rnn = nn.GRU(
                cnn_hidden_size, rnn_hidden_size, batch_first=False)
            self.rnn = weight_norm_rnn(self.rnn)
        else:
            self.rnn = nn.LSTM(
                cnn_hidden_size, rnn_hidden_size, batch_first=False)
        if skip:
            if rnn_type == "GRU":
                self.skiprnn = nn.GRU(
                    cnn_hidden_size, skip_hidden_size, batch_first=False)
                self.skiprnn = weight_norm_rnn(self.skiprnn)
            elif rnn_type == "SRU":
                self.skiprnn = nn.SRU(cnn_hidden_size, skip_hidden_size)
            else:
                self.skiprnn = nn.LSTM(
                    cnn_hidden_size, skip_hidden_size, batch_first=False)
            output_dim = rnn_hidden_size + \
                skip_hidden_size * skip + steps * 7
        else:
            output_dim = rnn_hidden_size + steps * 7
        self.output_network = nn.Sequential(
            nn.utils.parametrizations.weight_norm(
                nn.Linear(output_dim, rnn_hidden_size)),
            nn.ReLU(),
            nn.Dropout(self.odrop),
            nn.utils.parametrizations.weight_norm(
                nn.Linear(rnn_hidden_size, steps + 1))
        )
        if self.ar_window_size:
            self.ar_highway = nn.utils.parametrizations.weight_norm(
                nn.Linear(ar_window_size, steps + 1))
        self.init_mlp_weights()
//  [python]
def reset_parameters(self):
        # One problem with initialization from the uniform distribution is that
        # the distribution of the outputs has a variance that grows with the
        # number of inputs. It turns out that we can normalize the variance of
        # each neuron’s output to 1 by scaling its weight vector by the square
        # root of its fan-in (i.e. its number of inputs). Dropout further
        # increases the variance of each input, so we need to scale down std.
        # See A.3. in Gehring et al (2017): https://arxiv.org/pdf/1705.03122.
        std = math.sqrt((1 - self.dropout) / self.in_features)
        self.weight.data.normal_(mean=0, std=std)
        if self.bias is not None:
            self.bias.data.fill_(0)

        # Weight normalization is a reparameterization that decouples the
        # magnitude of a weight tensor from its direction. See Salimans and
        # Kingma (2016): https://arxiv.org/abs/1602.07868.
        if self.weight_norm:
            nn.utils.weight_norm(self)
//  [python]
def reset_parameters(self):
        # One problem with initialization from the uniform distribution is that
        # the distribution of the outputs has a variance that grows with the
        # number of inputs. It turns out that we can normalize the variance of
        # each neuron’s output to 1 by scaling its weight vector by the square
        # root of its fan-in (i.e. its number of inputs). Dropout further
        # increases the variance of each input, so we need to scale down std.
        # See A.3. in Gehring et al (2017): https://arxiv.org/pdf/1705.03122.
        std = math.sqrt((1 - self.dropout) / self.in_features)
        self.weight.data.normal_(mean=0, std=std)
        if self.bias is not None:
            self.bias.data.fill_(0)

        # Weight normalization is a reparameterization that decouples the
        # magnitude of a weight tensor from its direction. See Salimans and
        # Kingma (2016): https://arxiv.org/abs/1602.07868.
        if self.weight_norm:
            nn.utils.parametrizations.weight_norm(self)
//  [python]
def __init__(self, channel, dilation=1):
        super(ResStack, self).__init__()
        self.block = nn.Sequential(nn.LeakyReLU(0.2), nn.ReflectionPad1d(dilation), nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=dilation)), nn.LeakyReLU(0.2), nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)))
        self.shortcut = nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
//  [python]
def __init__(self, channel, dilation=1):
        super(ResStack, self).__init__()
        self.block = nn.Sequential(nn.LeakyReLU(0.2), nn.ReflectionPad1d(dilation), nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=dilation)), nn.LeakyReLU(0.2), nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)))
        self.shortcut = nn.utils.parametrizations.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
//  [python]
def apply_weight_norm(self):
        """
        Apply weight normalization module from all of the layers.
        """

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.utils.weight_norm(m)

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

        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.utils.parametrizations.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def __init__(self, period):
        super(PeriodicDiscriminator, self).__init__()
        self.period = period #入力波形を何分割するか
        self.kernel_size = 5
        self.stride = 3
        self.padding = 2

        self.conv2ds = nn.ModuleList([
            torch.nn.utils.weight_norm(nn.Conv2d(1, 32, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.weight_norm(nn.Conv2d(32, 128, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.weight_norm(nn.Conv2d(128, 512, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.weight_norm(nn.Conv2d(512, 1024, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.weight_norm(nn.Conv2d(1024, 1024, (self.kernel_size, 1), 1, padding=self.padding)),
        ])
        self.conv2d_post = torch.nn.utils.weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
//  [python]
def __init__(self, period):
        super(PeriodicDiscriminator, self).__init__()
        self.period = period #入力波形を何分割するか
        self.kernel_size = 5
        self.stride = 3
        self.padding = 2

        self.conv2ds = nn.ModuleList([
            torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(1, 32, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(32, 128, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(128, 512, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(512, 1024, (self.kernel_size, 1), (self.stride, 1), padding=(self.padding, 0))),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(1024, 1024, (self.kernel_size, 1), 1, padding=self.padding)),
        ])
        self.conv2d_post = torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
//  [python]
def __init__(self):
        super(Discriminator, self).__init__()
        self.conv1ds = nn.ModuleList([
            torch.nn.utils.weight_norm(nn.Conv1d(1, 16, 15, 1, padding=7)),
            torch.nn.utils.weight_norm(nn.Conv1d(16, 64, 41, 4, groups=4, padding=20)),
            torch.nn.utils.weight_norm(nn.Conv1d(64, 256, 41, 4, groups=16, padding=20)),
            torch.nn.utils.weight_norm(nn.Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
            torch.nn.utils.weight_norm(nn.Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
            torch.nn.utils.weight_norm(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
        ])
        self.conv1d_post = torch.nn.utils.weight_norm(nn.Conv1d(1024, 1, 3, 1, padding=1))
//  [python]
def __init__(self):
        super(Discriminator, self).__init__()
        self.conv1ds = nn.ModuleList([
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(1, 16, 15, 1, padding=7)),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(16, 64, 41, 4, groups=4, padding=20)),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(64, 256, 41, 4, groups=16, padding=20)),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
            torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
        ])
        self.conv1d_post = torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(1024, 1, 3, 1, padding=1))
//  [python]
def __init__(self, in_channels, out_channels, kernel_size, dilation,
                 final=False):
        super(CausalConvolutionBlock, self).__init__()
        
        # Computes left padding so that the applied convolutions are causal
        self.padding = (kernel_size - 1) * dilation
        padding = self.padding
        # First causal convolution
        self.conv1 = torch.nn.utils.weight_norm(torch.nn.Conv1d(
            in_channels, out_channels, kernel_size,
            padding=padding, dilation=dilation
        ))
        # The truncation makes the convolution causal
        self.chomp1 = Chomp1d(padding)
        self.dropout1 = nn.Dropout(0.1)

        # Second causal convolution
        self.conv2 = torch.nn.utils.weight_norm(torch.nn.Conv1d(
            out_channels, out_channels, kernel_size,
            padding=padding, dilation=dilation
        ))
        self.chomp2 = Chomp1d(padding)
        self.dropout2 = nn.Dropout(0.1)

        # Residual connection
        self.upordownsample = torch.nn.Conv1d(
            in_channels, out_channels, 1
        ) if in_channels != out_channels else None

        # Final activation function
        self.relu = torch.nn.ReLU() if final else None
//  [python]
def __init__(self, in_channels, out_channels, kernel_size, dilation,
                 final=False):
        super(CausalConvolutionBlock, self).__init__()
        
        # Computes left padding so that the applied convolutions are causal
        self.padding = (kernel_size - 1) * dilation
        padding = self.padding
        # First causal convolution
        self.conv1 = torch.nn.utils.parametrizations.weight_norm(torch.nn.Conv1d(
            in_channels, out_channels, kernel_size,
            padding=padding, dilation=dilation
        ))
        # The truncation makes the convolution causal
        self.chomp1 = Chomp1d(padding)
        self.dropout1 = nn.Dropout(0.1)

        # Second causal convolution
        self.conv2 = torch.nn.utils.parametrizations.weight_norm(torch.nn.Conv1d(
            out_channels, out_channels, kernel_size,
            padding=padding, dilation=dilation
        ))
        self.chomp2 = Chomp1d(padding)
        self.dropout2 = nn.Dropout(0.1)

        # Residual connection
        self.upordownsample = torch.nn.Conv1d(
            in_channels, out_channels, 1
        ) if in_channels != out_channels else None

        # Final activation function
        self.relu = torch.nn.ReLU() if final else None
//  [python]
def forward(self, x1):
        w1 = torch.nn.utils.weight_norm(torch.nn.Linear(3, 2))
        v1 = w1(x1)
        v2 = v1 + 3
        v3 = torch.clamp_min(v2, 0)
        v4 = torch.clamp_max(v3, 6)
        v5 = v4 / 6
        return v5
//  [python]
def forward(self, x1):
        w1 = torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(3, 2))
        v1 = w1(x1)
        v2 = v1 + 3
        v3 = torch.clamp_min(v2, 0)
        v4 = torch.clamp_max(v3, 6)
        v5 = v4 / 6
        return v5
//  [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 = nn.Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = ResBlock1 if resblock == '1' else ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(
                weight_norm(
                    nn.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 = nn.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 = nn.Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        resblock = ResBlock1 if resblock == '1' else ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(
                parametrizations.weight_norm(
                    nn.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 = nn.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, in_channels, mid_channels, num_blocks, mask_type, reverse_mask):
        super(CouplingLayer, self).__init__()

        # Save mask info
        self.mask_type = mask_type
        self.reverse_mask = reverse_mask

        if self.mask_type == MaskType.CHANNEL_WISE:
            in_channels //= 2

        # Build scale and translate network
        growth_rate, steps = 64, 5
        self.st_net = nn.Sequential(
            DenseNetLayer(inplanes=in_channels,
                          growth_rate=growth_rate,
                          steps=steps),
            Conv2d(in_channels + growth_rate*steps, 2*in_channels, 
                    kernel_size=3, stride=1, padding=1)
        )

        # Learnable scale for s
        self.rescale = nn.utils.weight_norm(Rescale(in_channels))
//  [python]
def __init__(self, in_channels, mid_channels, num_blocks, mask_type, reverse_mask):
        super(CouplingLayer, self).__init__()

        # Save mask info
        self.mask_type = mask_type
        self.reverse_mask = reverse_mask

        if self.mask_type == MaskType.CHANNEL_WISE:
            in_channels //= 2

        # Build scale and translate network
        growth_rate, steps = 64, 5
        self.st_net = nn.Sequential(
            DenseNetLayer(inplanes=in_channels,
                          growth_rate=growth_rate,
                          steps=steps),
            Conv2d(in_channels + growth_rate*steps, 2*in_channels, 
                    kernel_size=3, stride=1, padding=1)
        )

        # Learnable scale for s
        self.rescale = nn.utils.parametrizations.weight_norm(Rescale(in_channels))
//  [python]
def __init__(
        self,
        latent_size,
        output_dim = 1,
        norm_size=128,
        dims=[ 512, 512, 512, 512, 512, 512, 512, 512],
        dropout=[0, 1, 2, 3, 4, 5, 6, 7],
        dropout_prob=0.2,
        norm_layers=[0, 1, 2, 3, 4, 5, 6, 7],
        latent_in=[4],
        weight_norm=True,
        xyz_in_all=None,
        use_tanh=False,
        latent_dropout=False,
    ):
        super(Decoder, self).__init__()

        dims = [latent_size + norm_size] + dims + [output_dim]

        self.num_layers = len(dims)
        self.norm_layers = norm_layers
        self.latent_in = latent_in
        self.latent_dropout = latent_dropout
        if self.latent_dropout:
            self.lat_dp = nn.Dropout(0.2)

        self.xyz_in_all = xyz_in_all
        self.weight_norm = weight_norm
        self.norm_size = norm_size

        for l in range(0, self.num_layers - 1):
            if l + 1 in latent_in:
                out_dim = dims[l + 1] - dims[0]
            else:
                out_dim = dims[l + 1]
                if self.xyz_in_all and l != self.num_layers - 2:
                    out_dim -= norm_size

            if weight_norm and l in self.norm_layers:
                setattr(self, "lin" + str(l), nn.utils.weight_norm(nn.Linear(dims[l], out_dim)))
            else:
                setattr(self, "lin" + str(l), nn.Linear(dims[l], out_dim))

            if (not weight_norm) and self.norm_layers is not None and l in self.norm_layers:
                setattr(self, "bn" + str(l), nn.LayerNorm(out_dim))

        self.use_tanh = use_tanh
        if use_tanh:
            self.tanh = nn.Tanh()
        self.relu = nn.ReLU()

        self.dropout_prob = dropout_prob
        self.dropout = dropout
//  [python]
def __init__(
        self,
        latent_size,
        output_dim = 1,
        norm_size=128,
        dims=[ 512, 512, 512, 512, 512, 512, 512, 512],
        dropout=[0, 1, 2, 3, 4, 5, 6, 7],
        dropout_prob=0.2,
        norm_layers=[0, 1, 2, 3, 4, 5, 6, 7],
        latent_in=[4],
        weight_norm=True,
        xyz_in_all=None,
        use_tanh=False,
        latent_dropout=False,
    ):
        super(Decoder, self).__init__()

        dims = [latent_size + norm_size] + dims + [output_dim]

        self.num_layers = len(dims)
        self.norm_layers = norm_layers
        self.latent_in = latent_in
        self.latent_dropout = latent_dropout
        if self.latent_dropout:
            self.lat_dp = nn.Dropout(0.2)

        self.xyz_in_all = xyz_in_all
        self.weight_norm = weight_norm
        self.norm_size = norm_size

        for l in range(0, self.num_layers - 1):
            if l + 1 in latent_in:
                out_dim = dims[l + 1] - dims[0]
            else:
                out_dim = dims[l + 1]
                if self.xyz_in_all and l != self.num_layers - 2:
                    out_dim -= norm_size

            if weight_norm and l in self.norm_layers:
                setattr(self, "lin" + str(l), nn.utils.parametrizations.weight_norm(nn.Linear(dims[l], out_dim)))
            else:
                setattr(self, "lin" + str(l), nn.Linear(dims[l], out_dim))

            if (not weight_norm) and self.norm_layers is not None and l in self.norm_layers:
                setattr(self, "bn" + str(l), nn.LayerNorm(out_dim))

        self.use_tanh = use_tanh
        if use_tanh:
            self.tanh = nn.Tanh()
        self.relu = nn.ReLU()

        self.dropout_prob = dropout_prob
        self.dropout = dropout
//  [python]
def _gen_conv(
        self,
        in_ch,
        out_ch,
        strides=(2, 1),
        kernel_size=(5, 3),
        gain=math.sqrt(2),
        pad=(1, 1, 1, 2),
        use_weight_norm=False,
    ):
        pad = torch.nn.ReplicationPad2d(pad)
        conv = nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=strides, padding=0)

        if use_weight_norm:
            conv = torch.nn.utils.weight_norm(conv, name='weight')

        w = conv.weight
        k = w.size(1) * w.size(2) * w.size(3)
        conv.weight.data.normal_(0.0, gain / math.sqrt(k))
        nn.init.constant_(conv.bias, 0.01)
        return conv, pad
//  [python]
def _gen_conv(
        self,
        in_ch,
        out_ch,
        strides=(2, 1),
        kernel_size=(5, 3),
        gain=math.sqrt(2),
        pad=(1, 1, 1, 2),
        use_weight_norm=False,
    ):
        pad = torch.nn.ReplicationPad2d(pad)
        conv = nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=strides, padding=0)

        if use_weight_norm:
            conv = torch.nn.utils.parametrizations.weight_norm(conv, name='weight')

        w = conv.weight
        k = w.size(1) * w.size(2) * w.size(3)
        conv.weight.data.normal_(0.0, gain / math.sqrt(k))
        nn.init.constant_(conv.bias, 0.01)
        return conv, pad
//  [python]
def LinearActivation(
    d_input, d_output, bias=True,
    zero_bias_init=False,
    transposed=False,
    initializer=None,
    activation=None,
    activate=False,  # Apply activation as part of this module
    weight_norm=False,
    **kwargs,
):
    """ Returns a linear nn.Module with control over axes order, initialization, and activation """

    # Construct core module
    linear_cls = TransposedLinear if transposed else nn.Linear
    if activation == 'glu':
        d_output *= 2
    linear = linear_cls(d_input, d_output, bias=bias, **kwargs)

    # Initialize weight
    if initializer is not None:
        get_initializer(initializer, activation)(linear.weight)

    # Initialize bias
    if bias and zero_bias_init:
        nn.init.zeros_(linear.bias)

    # Weight norm
    if weight_norm:
        linear = nn.utils.weight_norm(linear)

    if activate and activation is not None:
        activation = Activation(activation, dim=-2 if transposed else -1)
        linear = nn.Sequential(linear, activation)
    return linear