# 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
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self.bidirectional = bidirectional
        self.lstm = False
        num_directions = 2 if bidirectional else 1

        self.all_cells = []
        for layer in range(num_layers):
            for direction in range(num_directions):
                layer_input_size = input_size if layer == 0 else hidden_size * num_directions

                cell = self.Cell(layer_input_size, hidden_size, self.bias, p=dropout, **kwargs)
                self.all_cells.append(cell)
                self.add_module('cell%d' % (layer * num_directions + direction), cell)

    def reset_parameters(self):
        for cell in self.all_cells:
            cell.reset_parameters()

    def reset_noise(self, batch_size):
        for cell in self.all_cells:
            cell.reset_noise(batch_size)

    def forward(self, input, skip_connect, mask=None, hx=None):
        batch_size = input.size(0) if self.batch_first else input.size(1)
        if hx is None:
            num_directions = 2 if self.bidirectional else 1

# 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

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self.bidirectional = bidirectional
        self.lstm = False
        num_directions = 2 if bidirectional else 1

        self.all_cells = []
        for layer in range(num_layers):
            for direction in range(num_directions):
                layer_input_size = input_size if layer == 0 else hidden_size * num_directions

                cell = self.Cell(layer_input_size, hidden_size, self.bias, p=dropout, **kwargs)
                self.all_cells.append(cell)
                self.add_module('cell%d' % (layer * num_directions + direction), cell)

    def reset_parameters(self):
        for cell in self.all_cells:
            cell.reset_parameters()

    def reset_noise(self, batch_size):
        for cell in self.all_cells:
            cell.reset_noise(batch_size)

    def forward(self, input, skip_connect, mask=None, hx=None):
        batch_size = input.size(0) if self.batch_first else input.size(1)
        if hx is None:
            num_directions = 2 if self.bidirectional else 1

            hx = input.new_zeros(self.num_layers * num_directions, batch_size, self.hidden_size)