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

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

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(
                gin_channels, 2 * hidden_channels * n_layers, 1
            )
            self.cond_layer = torch.nn.utils.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 __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
        sigmoid_scale=False,
        n_sqz=2,
    ):
        super().__init__()
        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.sigmoid_scale = sigmoid_scale
        self.n_sqz = n_sqz

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

        self.wn = WN(
            hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout, gin_channels
        )
        self.wn_pitch = WNP(
            hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout, 1, n_sqz
        )
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
        sigmoid_scale=False,
        n_sqz=2,
    ):
        super().__init__()
        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.sigmoid_scale = sigmoid_scale
        self.n_sqz = n_sqz

        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start)
        self.start = start
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        self.wn = WN(
            hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout, gin_channels
        )
        self.wn_pitch = WNP(
            hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout, 1, n_sqz
        )
//  [python]
def __init__(
        self, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0, n_sqz=2,
    ):
        super(WNP, self).__init__()
        assert kernel_size % 2 == 1
        assert hidden_channels % 2 == 0

        self.hidden_channels = hidden_channels
        self.n_layers = n_layers

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

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

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

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

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

        self.hidden_channels = hidden_channels
        self.n_layers = n_layers

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

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers // self.n_sqz, 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 __init__(self, slayer, inFeatures, outFeatures, weightScale=100, 
                 preHookFx = lambda x: utils.quantize(x, step=2), weightNorm=False, 
                 delay=False, maxDelay=62, countLog=False):
        super(denseBlock, self).__init__()
        self.slayer = slayer
        self.weightNorm = weightNorm
        if weightNorm is True:
            self.weightOp = torch.nn.utils.weight_norm(slayer.dense(inFeatures, outFeatures, weightScale, preHookFx), name='weight')
        else:
            self.weightOp = slayer.dense(inFeatures, outFeatures, weightScale, preHookFx)
        self.delayOp  = slayer.delay(outFeatures) if delay is True else None
        self.countLog = countLog
        self.gradLog = True
        self.maxDelay = maxDelay

        self.paramsDict = {
            'inFeatures'  : inFeatures,
            'outFeatures' : outFeatures,
        }
//  [python]
def __init__(self, slayer, inFeatures, outFeatures, weightScale=100, 
                 preHookFx = lambda x: utils.quantize(x, step=2), weightNorm=False, 
                 delay=False, maxDelay=62, countLog=False):
        super(denseBlock, self).__init__()
        self.slayer = slayer
        self.weightNorm = weightNorm
        if weightNorm is True:
            self.weightOp = torch.nn.utils.parametrizations.weight_norm(slayer.dense(inFeatures, outFeatures, weightScale, preHookFx), name='weight')
        else:
            self.weightOp = slayer.dense(inFeatures, outFeatures, weightScale, preHookFx)
        self.delayOp  = slayer.delay(outFeatures) if delay is True else None
        self.countLog = countLog
        self.gradLog = True
        self.maxDelay = maxDelay

        self.paramsDict = {
            'inFeatures'  : inFeatures,
            'outFeatures' : outFeatures,
        }
//  [python]
def __init__(self, n_in_channels, n_layers, n_channels, kernel_size=3):
        super(WN, self).__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()

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

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

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

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

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

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

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

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, in_channels, cin_channels, di_cycle, pos_group, n_channels, n_layers, kernel_size=3):
        super().__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.pos_group = pos_group

        self.in_layers = torch.nn.ModuleList()
        self.time_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()

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

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        out_channels = in_channels * 2
        end = torch.nn.Conv1d(n_channels, out_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end
        cond_layer = torch.nn.Conv1d(cin_channels, 2*n_channels*n_layers, 1)
        self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
        
        if pos_group > 1:
            pos_emb = torch.nn.Embedding(pos_group, 2*n_channels*n_layers)
            self.pos_emb = torch.nn.utils.weight_norm(pos_emb, name='weight')

        for i in range(n_layers):
            dilation = 2 ** (i % di_cycle)
            padding = int((kernel_size*dilation - dilation)) // 2

            in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(
                res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, in_channels, cin_channels, di_cycle, pos_group, n_channels, n_layers, kernel_size=3):
        super().__init__()
        assert(kernel_size % 2 == 1)
        assert(n_channels % 2 == 0)
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.pos_group = pos_group

        self.in_layers = torch.nn.ModuleList()
        self.time_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()

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

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        out_channels = in_channels * 2
        end = torch.nn.Conv1d(n_channels, out_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end
        cond_layer = torch.nn.Conv1d(cin_channels, 2*n_channels*n_layers, 1)
        self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')
        
        if pos_group > 1:
            pos_emb = torch.nn.Embedding(pos_group, 2*n_channels*n_layers)
            self.pos_emb = torch.nn.utils.parametrizations.weight_norm(pos_emb, name='weight')

        for i in range(n_layers):
            dilation = 2 ** (i % di_cycle)
            padding = int((kernel_size*dilation - dilation)) // 2

            in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2*n_channels
            else:
                res_skip_channels = n_channels
            res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(
                res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, n_sqz=1):
        super(NonCausalWaveNet, self).__init__()
        assert(kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        # self.gin_channels = gin_channels
        self.p_dropout = p_dropout

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

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

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

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

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

        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, input_dim: int, layers: Sequence[Sequence[int]], direction: str, do_weight_norm: bool=True, dropout: float=0.0) ->None:
        super().__init__()
        self.dropout = dropout
        self._convolutions = torch.nn.ModuleList()
        last_dim = input_dim
        for k, layer in enumerate(layers):
            if len(layer) == 2:
                conv = torch.nn.Conv1d(last_dim, layer[1] * 2, layer[0], stride=1, padding=layer[0] - 1, bias=True)
            elif len(layer) == 3:
                assert layer[0] == 2, 'only support kernel = 2 for now'
                conv = torch.nn.Conv1d(last_dim, layer[1] * 2, layer[0], stride=1, padding=layer[2], dilation=layer[2], bias=True)
            else:
                raise ValueError('each layer must have length 2 or 3')
            if k == 0:
                conv_dropout = dropout
            else:
                conv_dropout = 0.0
            std = math.sqrt(4 * (1.0 - conv_dropout) / (layer[0] * last_dim))
            conv.weight.data.normal_(0, std=std)
            conv.bias.data.zero_()
            if do_weight_norm:
                conv = torch.nn.utils.weight_norm(conv, name='weight', dim=0)
            self._convolutions.append(conv)
            last_dim = layer[1]
        assert last_dim == input_dim
        if direction not in ('forward', 'backward'):
            raise ConfigurationError(f'invalid direction: {direction}')
        self._direction = direction
//  [python]
def __init__(self, input_dim: int, layers: Sequence[Sequence[int]], direction: str, do_weight_norm: bool=True, dropout: float=0.0) ->None:
        super().__init__()
        self.dropout = dropout
        self._convolutions = torch.nn.ModuleList()
        last_dim = input_dim
        for k, layer in enumerate(layers):
            if len(layer) == 2:
                conv = torch.nn.Conv1d(last_dim, layer[1] * 2, layer[0], stride=1, padding=layer[0] - 1, bias=True)
            elif len(layer) == 3:
                assert layer[0] == 2, 'only support kernel = 2 for now'
                conv = torch.nn.Conv1d(last_dim, layer[1] * 2, layer[0], stride=1, padding=layer[2], dilation=layer[2], bias=True)
            else:
                raise ValueError('each layer must have length 2 or 3')
            if k == 0:
                conv_dropout = dropout
            else:
                conv_dropout = 0.0
            std = math.sqrt(4 * (1.0 - conv_dropout) / (layer[0] * last_dim))
            conv.weight.data.normal_(0, std=std)
            conv.bias.data.zero_()
            if do_weight_norm:
                conv = torch.nn.utils.parametrizations.weight_norm(conv, name='weight', dim=0)
            self._convolutions.append(conv)
            last_dim = layer[1]
        assert last_dim == input_dim
        if direction not in ('forward', 'backward'):
            raise ConfigurationError(f'invalid direction: {direction}')
        self._direction = direction
//  [python]
def weight_norm_modules(module, name = 'weight', dim = 0):
    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
        module.weight_norm()
        return module
    else:
        return weight_norm(module,name,dim)
//  [python]
def weight_norm_modules(module, name = 'weight', dim = 0):
    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
        module.weight_norm()
        return module
    else:
        return torch.nn.utils.parametrizations.weight_norm(module,name,dim)
//  [python]
def __init__(self, slayer, inChannels, outChannels, kernelSize, stride=1, padding=0, dilation=1, groups=1, weightScale=100, 
                 preHookFx = lambda x: utils.quantize(x, step=2), weightNorm=False, 
                 delay=False, maxDelay=62, countLog=False):
        super(convBlock, self).__init__()
        self.slayer = slayer
        self.weightNorm = weightNorm
        if weightNorm is True:
            self.weightOp = torch.nn.utils.weight_norm(
                slayer.conv(inChannels, outChannels, kernelSize, stride, padding, dilation, groups, weightScale, preHookFx), 
                name='weight',
            )
        else:
            self.weightOp = slayer.conv(inChannels, outChannels, kernelSize, stride, padding, dilation, groups, weightScale, preHookFx)
        # only channel wise delay is supported for conv layer
        # for neuron wise delay, one will need to write a custom block as it would require the spatial dimension as well
        self.delayOp  = slayer.delay(outChannels) if delay is True else None    
        self.countLog = countLog
        self.gradLog = True
        self.maxDelay = maxDelay

        self.paramsDict = {
            'inChannels'  : inChannels,
            'outChannels' : outChannels,
            'kernelSize'  : kernelSize,
            'stride'      : stride,
            'padding'     : padding,
            'dilation'    : dilation,
            'groups'      : groups,
        }
//  [python]
def __init__(self, slayer, inChannels, outChannels, kernelSize, stride=1, padding=0, dilation=1, groups=1, weightScale=100, 
                 preHookFx = lambda x: utils.quantize(x, step=2), weightNorm=False, 
                 delay=False, maxDelay=62, countLog=False):
        super(convBlock, self).__init__()
        self.slayer = slayer
        self.weightNorm = weightNorm
        if weightNorm is True:
            self.weightOp = torch.nn.utils.parametrizations.weight_norm(
                slayer.conv(inChannels, outChannels, kernelSize, stride, padding, dilation, groups, weightScale, preHookFx), 
                name='weight',
            )
        else:
            self.weightOp = slayer.conv(inChannels, outChannels, kernelSize, stride, padding, dilation, groups, weightScale, preHookFx)
        # only channel wise delay is supported for conv layer
        # for neuron wise delay, one will need to write a custom block as it would require the spatial dimension as well
        self.delayOp  = slayer.delay(outChannels) if delay is True else None    
        self.countLog = countLog
        self.gradLog = True
        self.maxDelay = maxDelay

        self.paramsDict = {
            'inChannels'  : inChannels,
            'outChannels' : outChannels,
            'kernelSize'  : kernelSize,
            'stride'      : stride,
            'padding'     : padding,
            'dilation'    : dilation,
            'groups'      : groups,
        }
//  [python]
def __init__(self,num_features, length_hidden):
        super().__init__()
        # self.net = nn.Linear(num_features, num_features)
        self.fc1 = nn.Linear(num_features, int(length_hidden*num_features))
        self.fc2 = nn.Linear(int(length_hidden*num_features), int(length_hidden*num_features))
        self.fc3 = nn.Linear(int(length_hidden*num_features), num_features)
        # self.fc4 = nn.Linear(int(length_hidden * num_features), int(length_hidden*num_features))
        # self.fc5 = nn.Linear(int(length_hidden * num_features), num_features)
        # self.fc6 = nn.Linear(int(length_hidden * num_features), num_features)
        self.rescale = nn.utils.weight_norm(Rescale(num_features))
//  [python]
def __init__(self,num_features, length_hidden):
        super().__init__()
        # self.net = nn.Linear(num_features, num_features)
        self.fc1 = nn.Linear(num_features, int(length_hidden*num_features))
        self.fc2 = nn.Linear(int(length_hidden*num_features), int(length_hidden*num_features))
        self.fc3 = nn.Linear(int(length_hidden*num_features), num_features)
        # self.fc4 = nn.Linear(int(length_hidden * num_features), int(length_hidden*num_features))
        # self.fc5 = nn.Linear(int(length_hidden * num_features), num_features)
        # self.fc6 = nn.Linear(int(length_hidden * num_features), num_features)
        self.rescale = nn.utils.parametrizations.weight_norm(Rescale(num_features))
//  [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.Conv1d):
                torch.nn.utils.weight_norm(m)
                logging.info(f"Weight norm is applied to {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.Conv1d):
                torch.nn.utils.parametrizations.weight_norm(m)
                logging.info(f"Weight norm is applied to {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.Conv1d) \
                or isinstance(m, torch.nn.ConvTranspose2d):
                torch.nn.utils.weight_norm(m)
                logging.info(f"Weight norm is applied to {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.Conv1d) \
                or isinstance(m, torch.nn.ConvTranspose2d):
                torch.nn.utils.parametrizations.weight_norm(m)
                logging.info(f"Weight norm is applied to {m}.")

        self.apply(_apply_weight_norm)
//  [python]
def loadModel(self, fname):
        '''
        This function loads the network from a perviously saved hdf5 file using ``genModel``.

        Usage:

        .. code-block:: python
            
            net = Network(netParams)
            net.loadModel(path_of_model)
        '''
        # only the layer weights and delays shall be loaded
        h = h5py.File(fname, 'r')

        # one more layer for input layer in the hdf5 file
        assert len(h['layer']) == len(self.blocks) + 1, 'The number of layers in the network does not match with the number of layers in the file {}. Expected {}, found {}'.format(fname, len(self.blocks) + 1, len(h['layer']))

        for i, block in enumerate(self.blocks):
            idxKey = '{}'.format(i+1)
            blockTypeStr = block.__class__.__name__[:-5]
            layerTypeStr = h['layer'][idxKey]['type'][()][0].decode('utf-8')
            assert layerTypeStr == blockTypeStr, 'The layer typestring do not match. Found {} in network and {} in file.'.format(blockTypeStr, layerTypeStr)

            if block.weightOp is not None:
                if self.weightNorm is True and layerTypeStr != 'pool':
                    torch.nn.utils.remove_weight_norm(block.weightOp, name='weight')
                block.weightOp.weight.data = torch.FloatTensor(h['layer'][idxKey]['weight'][()]).reshape(block.weightOp.weight.shape).to(block.weightOp.weight.device)
                if self.weightNorm is True and layerTypeStr != 'pool':
                    block.weightOp = torch.nn.utils.weight_norm(block.weightOp, name='weight')

            if block.delayOp is not None:
                block.delayOp.delay.data = torch.FloatTensor(h['layer'][idxKey]['delay'][()]).reshape(block.delayOp.delay.shape).to(block.delayOp.delay.device)
//  [python]
def loadModel(self, fname):
        '''
        This function loads the network from a perviously saved hdf5 file using ``genModel``.

        Usage:

        .. code-block:: python
            
            net = Network(netParams)
            net.loadModel(path_of_model)
        '''
        # only the layer weights and delays shall be loaded
        h = h5py.File(fname, 'r')

        # one more layer for input layer in the hdf5 file
        assert len(h['layer']) == len(self.blocks) + 1, 'The number of layers in the network does not match with the number of layers in the file {}. Expected {}, found {}'.format(fname, len(self.blocks) + 1, len(h['layer']))

        for i, block in enumerate(self.blocks):
            idxKey = '{}'.format(i+1)
            blockTypeStr = block.__class__.__name__[:-5]
            layerTypeStr = h['layer'][idxKey]['type'][()][0].decode('utf-8')
            assert layerTypeStr == blockTypeStr, 'The layer typestring do not match. Found {} in network and {} in file.'.format(blockTypeStr, layerTypeStr)

            if block.weightOp is not None:
                if self.weightNorm is True and layerTypeStr != 'pool':
                    torch.nn.utils.parametrize.remove_weight_norm(block.weightOp, name='weight')
                block.weightOp.weight.data = torch.FloatTensor(h['layer'][idxKey]['weight'][()]).reshape(block.weightOp.weight.shape).to(block.weightOp.weight.device)
                if self.weightNorm is True and layerTypeStr != 'pool':
                    block.weightOp = torch.nn.utils.parametrizations.weight_norm(block.weightOp, name='weight')

            if block.delayOp is not None:
                block.delayOp.delay.data = torch.FloatTensor(h['layer'][idxKey]['delay'][()]).reshape(block.delayOp.delay.shape).to(block.delayOp.delay.device)
//  [python]
def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.utils.weight_norm(torch.nn.Conv2d(3, 32, 3), name="weight")
        self.conv2 = torch.nn.Conv2d(32, 2, 1, stride=1, padding=0)
        self.conv3 = torch.nn.Conv2d(2, 1, 12, stride=1, padding=0)
//  [python]
def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.utils.parametrizations.weight_norm(torch.nn.Conv2d(3, 32, 3), name="weight")
        self.conv2 = torch.nn.Conv2d(32, 2, 1, stride=1, padding=0)
        self.conv3 = torch.nn.Conv2d(2, 1, 12, stride=1, padding=0)
//  [python]
def _apply_weight_norm(self):
        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d):
                torch.nn.utils.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def _apply_weight_norm(self):
        def _apply_weight_norm(m):
            if isinstance(m, torch.nn.Conv1d):
                torch.nn.utils.parametrizations.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def __init__(self, num_class=64, inc=4096, temp=0.1,
                 cosine=False,dropout=False):
        super(Predictor_2, self).__init__()

        if cosine:
            self.fc =  torch.nn.utils.weight_norm(nn.Linear(inc, num_class, bias=False), name='weight')
        else:
            self.fc = torch.nn.utils.weight_norm(nn.Linear(inc, num_class, bias=True), name='weight')

        # if cosine:
        #     self.fc = nn.Linear(inc, num_class, bias=False)
        # else:
        #     self.fc = nn.Linear(inc, num_class,bias=True)

        self.num_class = num_class
        self.temp = temp
        self.cosine = cosine
        self.dropout = dropout
//  [python]
def __init__(self, num_class=64, inc=4096, temp=0.1,
                 cosine=False,dropout=False):
        super(Predictor_2, self).__init__()

        if cosine:
            self.fc =  torch.nn.utils.parametrizations.weight_norm(nn.Linear(inc, num_class, bias=False), dim=0)
        else:
            self.fc = torch.nn.utils.parametrizations.weight_norm(nn.Linear(inc, num_class, bias=True), dim=0)

        # if cosine:
        #     self.fc = nn.Linear(inc, num_class, bias=False)
        # else:
        #     self.fc = nn.Linear(inc, num_class,bias=True)

        self.num_class = num_class
        self.temp = temp
        self.cosine = cosine
        self.dropout = dropout
//  [python]
def __init__(self, in_channels, out_channels, kernel_size, padding, bias=True):
		super(WNConv2d, self).__init__()
		self.conv = nn.utils.weight_norm(
			nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=bias))
//  [python]
def __init__(self, in_channels, out_channels, kernel_size, padding, bias=True):
		super(WNConv2d, self).__init__()
		self.conv = nn.utils.parametrizations.weight_norm(
			nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=bias))
//  [python]
def __init__(self, input_channel=80, hu=512, ku=[16, 16, 4, 4], kr=[3, 7, 11], Dr=[1, 3, 5]):
        super(Generator, self).__init__()
        self.input = nn.Sequential(
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(input_channel, hu, kernel_size=7))
        )

        generator = []
        
        for k in ku:
            inp = hu
            out = int(inp/2)
            generator += [
                nn.LeakyReLU(0.2),
                nn.utils.weight_norm(nn.ConvTranspose1d(inp, out, k, k//2)),
                MRF(kr, out, Dr)
            ]
            hu = out
        self.generator = nn.Sequential(*generator)

        self.output = nn.Sequential(
            nn.LeakyReLU(0.2),
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(hu, 1, kernel_size=7, stride=1)),
            nn.Tanh()

        )
//  [python]
def __init__(self, input_channel=80, hu=512, ku=[16, 16, 4, 4], kr=[3, 7, 11], Dr=[1, 3, 5]):
        super(Generator, self).__init__()
        self.input = nn.Sequential(
            nn.ReflectionPad1d(3),
            nn.utils.parametrizations.weight_norm(nn.Conv1d(input_channel, hu, kernel_size=7))
        )

        generator = []
        
        for k in ku:
            inp = hu
            out = int(inp/2)
            generator += [
                nn.LeakyReLU(0.2),
                nn.utils.parametrizations.weight_norm(nn.ConvTranspose1d(inp, out, k, k//2)),
                MRF(kr, out, Dr)
            ]
            hu = out
        self.generator = nn.Sequential(*generator)

        self.output = nn.Sequential(
            nn.LeakyReLU(0.2),
            nn.ReflectionPad1d(3),
            nn.utils.parametrizations.weight_norm(nn.Conv1d(hu, 1, kernel_size=7, stride=1)),
            nn.Tanh()

        )
//  [python]
def __init__(self, h):
        super(Generator, self).__init__()
        self.h = h
        self.num_kernels = len(h.resblock_kernel_sizes)
        self.num_upsamples = len(h.upsample_rates)
        self.m_source = SourceModuleHnNSF(
            sampling_rate=h.sampling_rate,
            harmonic_num=8
        )
        self.noise_convs = nn.ModuleList()
        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
        resblock = ResBlock1 if h.resblock == '1' else ResBlock2

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

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

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

        self.conv_post = parametrizations.weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)
        self.upp = int(np.prod(h.upsample_rates))
//  [python]
def test_weight_norm(self):
        input = Variable(torch.randn(3, 5))
        m = nn.Linear(5, 7)
        expected_output = m(input)

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

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

        # test with dim=1
        m = torch.nn.utils.weight_norm(m, dim=1)
        self.assertEqual(m.weight_v.size(), m.weight.size())
        self.assertEqual(m.weight_g.size(), (1, 5))
        self.assertEqual(m(input), expected_output)
//  [python]
def test_weight_norm(self):
        input = Variable(torch.randn(3, 5))
        m = nn.Linear(5, 7)
        expected_output = m(input)

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

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

        # test with dim=1
        m = torch.nn.utils.parametrizations.weight_norm(m, dim=1)
        self.assertEqual(m.weight_v.size(), m.weight.size())
        self.assertEqual(m.weight_g.size(), (1, 5))
        self.assertEqual(m(input), expected_output)
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 num_layers,
                 c_in_channels=0,
                 dropout_p=0,
                 sigmoid_scale=False):
        super().__init__()
        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.sigmoid_scale = sigmoid_scale
        # input layer
        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
        start = torch.nn.utils.weight_norm(start)
        self.start = start
        # output layer
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end
        # coupling layers
        self.wn = WN(in_channels, hidden_channels, kernel_size, dilation_rate,
                     num_layers, c_in_channels, dropout_p)
//  [python]
def __init__(self,
                 in_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 num_layers,
                 c_in_channels=0,
                 dropout_p=0,
                 sigmoid_scale=False):
        super().__init__()
        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.sigmoid_scale = sigmoid_scale
        # input layer
        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start)
        self.start = start
        # output layer
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end
        # coupling layers
        self.wn = WN(in_channels, hidden_channels, kernel_size, dilation_rate,
                     num_layers, c_in_channels, dropout_p)
//  [python]
def __init__(self, in_channels, num_channels, num_blocks, num_components, drop_prob, use_attn=True, aux_channels=None):
		super(NN, self).__init__()
		self.k = num_components  # k = number of mixture components
		self.in_conv = WNConv2d(in_channels, num_channels, kernel_size=3, padding=1)
		self.mid_convs = nn.ModuleList([ConvAttnBlock(num_channels, drop_prob, use_attn, aux_channels)
										for _ in range(num_blocks)])
		self.out_conv = WNConv2d(num_channels, in_channels * (2 + 3 * self.k),
								 kernel_size=3, padding=1)
		self.rescale = weight_norm(Rescale(in_channels))
//  [python]
def __init__(self, in_channels, num_channels, num_blocks, num_components, drop_prob, use_attn=True, aux_channels=None):
		super(NN, self).__init__()
		self.k = num_components  # k = number of mixture components
		self.in_conv = WNConv2d(in_channels, num_channels, kernel_size=3, padding=1)
		self.mid_convs = nn.ModuleList([ConvAttnBlock(num_channels, drop_prob, use_attn, aux_channels)
										for _ in range(num_blocks)])
		self.out_conv = WNConv2d(num_channels, in_channels * (2 + 3 * self.k),
								 kernel_size=3, padding=1)
		self.rescale = nn.utils.parametrizations.weight_norm(Rescale(in_channels))
//  [python]
def init_layers(self) -> None:
        self.fc = nn.Linear(self.in_channels, self.num_classes, bias=False)
        self.fc = nn.utils.weight_norm(self.fc, name='weight', dim=0)
//  [python]
def init_layers(self) -> None:
        self.fc = nn.Linear(self.in_channels, self.num_classes, bias=False)
        self.fc = nn.utils.parametrizations.weight_norm(self.fc, name='weight', dim=0)
//  [python]
def __init__(self, d_model, num_heads=4, drop_prob=0.):
		super(GatedAttn, self).__init__()
		self.d_model = d_model
		self.num_heads = num_heads
		self.drop_prob = drop_prob
		self.in_proj = weight_norm(nn.Linear(d_model, 3 * d_model, bias=False))
		self.gate = weight_norm(nn.Linear(d_model, 2 * d_model))
//  [python]
def __init__(self, d_model, num_heads=4, drop_prob=0.):
		super(GatedAttn, self).__init__()
		self.d_model = d_model
		self.num_heads = num_heads
		self.drop_prob = drop_prob
		self.in_proj = nn.utils.parametrizations.weight_norm(nn.Linear(d_model, 3 * d_model, bias=False))
		self.gate = nn.utils.parametrizations.weight_norm(nn.Linear(d_model, 2 * d_model))