# Path: neuronlp2/nn/_functions/skipconnect_rnn.py
# def SkipConnectRNNReLUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None, noise_skip=None):
# def SkipConnectRNNTanhCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectLSTMCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectFastLSTMCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectGRUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectFastGRUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectRecurrent(reverse=False):
#     def forward(input, skip_connect, hidden, cell, mask):
# def StackedRNN(inners, num_layers, lstm=False):
#     def reverse_skip_connection(skip_connect):
#     def forward(input, skip_connect, hidden, cells, mask):
# def AutogradSkipConnectRNN(num_layers=1, batch_first=False, bidirectional=False, lstm=False):
#     def forward(input, skip_connect, cells, hidden, mask):
# def SkipConnectStep():
#     def forward(input, hidden, hidden_skip, cell, mask):
# def StackedStep(layer, num_layers, lstm=False):
#     def forward(input, hidden, hidden_skip, cells, mask):
# def AutogradSkipConnectStep(num_layers=1, lstm=False):
#     def forward(input, cells, hidden, hidden_skip, mask):
# 
# Path: neuronlp2/nn/variational_rnn.py
# class VarRNNCellBase(nn.Module):
#     def __repr__(self):
#         s = '{name}({model_dim}, {hidden_size}'
#         if 'bias' in self.__dict__ and self.bias is not True:
#             s += ', bias={bias}'
#         if 'nonlinearity' in self.__dict__ and self.nonlinearity != "tanh":
#             s += ', nonlinearity={nonlinearity}'
#         s += ')'
#         return s.format(name=self.__class__.__name__, **self.__dict__)
# 
#     def reset_noise(self, batch_size):
#         """
#         Should be overriden by all subclasses.
#         Args:
#             batch_size: (int) batch size of input.
#         """
#         raise NotImplementedError


import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from neuronlp2.nn._functions import skipconnect_rnn as rnn_F
from neuronlp2.nn.variational_rnn import VarRNNCellBase
        if self.bias:
            nn.init.constant_(self.bias_hh, 0.)
            nn.init.constant_(self.bias_ih, 0.)

    def reset_noise(self, batch_size):
        if self.training:
            if self.p_in:
                noise = self.weight_ih.new_empty(batch_size, self.input_size)
                self.noise_in = noise.bernoulli_(1.0 - self.p_in) / (1.0 - self.p_in)
            else:
                self.noise_in = None

            if self.p_hidden:
                noise = self.weight_hh.new_empty(batch_size, self.hidden_size * 2)
                self.noise_hidden = noise.bernoulli_(1.0 - self.p_hidden) / (1.0 - self.p_hidden)
            else:
                self.noise_hidden = None
        else:
            self.noise_in = None
            self.noise_hidden = None

    def forward(self, input, hx, hs):
        return rnn_F.SkipConnectFastGRUCell(
            input, hx, hs,
            self.weight_ih, self.weight_hh,
            self.bias_ih, self.bias_hh,
            self.noise_in, self.noise_hidden,
        )



# Path: neuronlp2/nn/_functions/skipconnect_rnn.py
# def SkipConnectRNNReLUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None, noise_skip=None):
# def SkipConnectRNNTanhCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectLSTMCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectFastLSTMCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectGRUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectFastGRUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def SkipConnectRecurrent(reverse=False):
#     def forward(input, skip_connect, hidden, cell, mask):
# def StackedRNN(inners, num_layers, lstm=False):
#     def reverse_skip_connection(skip_connect):
#     def forward(input, skip_connect, hidden, cells, mask):
# def AutogradSkipConnectRNN(num_layers=1, batch_first=False, bidirectional=False, lstm=False):
#     def forward(input, skip_connect, cells, hidden, mask):
# def SkipConnectStep():
#     def forward(input, hidden, hidden_skip, cell, mask):
# def StackedStep(layer, num_layers, lstm=False):
#     def forward(input, hidden, hidden_skip, cells, mask):
# def AutogradSkipConnectStep(num_layers=1, lstm=False):
#     def forward(input, cells, hidden, hidden_skip, mask):
# 
# Path: neuronlp2/nn/variational_rnn.py
# class VarRNNCellBase(nn.Module):
#     def __repr__(self):
#         s = '{name}({model_dim}, {hidden_size}'
#         if 'bias' in self.__dict__ and self.bias is not True:
#             s += ', bias={bias}'
#         if 'nonlinearity' in self.__dict__ and self.nonlinearity != "tanh":
#             s += ', nonlinearity={nonlinearity}'
#         s += ')'
#         return s.format(name=self.__class__.__name__, **self.__dict__)
# 
#     def reset_noise(self, batch_size):
#         """
#         Should be overriden by all subclasses.
#         Args:
#             batch_size: (int) batch size of input.
#         """
#         raise NotImplementedError


# Path: neuronlp2/nn/skip_rnn.py
import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from neuronlp2.nn._functions import skipconnect_rnn as rnn_F
from neuronlp2.nn.variational_rnn import VarRNNCellBase

        if self.bias:
            nn.init.constant_(self.bias_hh, 0.)
            nn.init.constant_(self.bias_ih, 0.)

    def reset_noise(self, batch_size):
        if self.training:
            if self.p_in:
                noise = self.weight_ih.new_empty(batch_size, self.input_size)
                self.noise_in = noise.bernoulli_(1.0 - self.p_in) / (1.0 - self.p_in)
            else:
                self.noise_in = None

            if self.p_hidden:
                noise = self.weight_hh.new_empty(batch_size, self.hidden_size * 2)
                self.noise_hidden = noise.bernoulli_(1.0 - self.p_hidden) / (1.0 - self.p_hidden)
            else:
                self.noise_hidden = None
        else:
            self.noise_in = None
            self.noise_hidden = None

    def forward(self, input, hx, hs):
        return rnn_F.SkipConnectFastGRUCell(
            input, hx, hs,
            self.weight_ih, self.weight_hh,
            self.bias_ih, self.bias_hh,
            self.noise_in, self.noise_hidden,
        )



class SkipConnectGRUCell(VarRNNCellBase):