# Path: neuronlp2/nn/_functions/variational_rnn.py
# def VarRNNReLUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarRNNTanhCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarLSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarFastLSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarGRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarFastGRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarRecurrent(reverse=False):
#     def forward(input, hidden, cell, mask):
# def StackedRNN(inners, num_layers, lstm=False):
#     def forward(input, hidden, cells, mask):
# def AutogradVarRNN(num_layers=1, batch_first=False, bidirectional=False, lstm=False):
#     def forward(input, cells, hidden, mask):
# def VarRNNStep():
#     def forward(input, hidden, cell, mask):
# def StackedStep(layer, num_layers, lstm=False):
#     def forward(input, hidden, cells, mask):
# def AutogradVarRNNStep(num_layers=1, lstm=False):
#     def forward(input, cells, hidden, mask):


import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from neuronlp2.nn._functions import variational_rnn as rnn_F
__author__ = 'max'



class VarRNNBase(nn.Module):
    def __init__(self, Cell, input_size, hidden_size,
                 num_layers=1, bias=True, batch_first=False,
                 dropout=(0, 0), bidirectional=False, **kwargs):

        super(VarRNNBase, self).__init__()
        self.Cell = Cell
        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):

# Path: neuronlp2/nn/_functions/variational_rnn.py
# def VarRNNReLUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarRNNTanhCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarLSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarFastLSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarGRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarFastGRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
# def VarRecurrent(reverse=False):
#     def forward(input, hidden, cell, mask):
# def StackedRNN(inners, num_layers, lstm=False):
#     def forward(input, hidden, cells, mask):
# def AutogradVarRNN(num_layers=1, batch_first=False, bidirectional=False, lstm=False):
#     def forward(input, cells, hidden, mask):
# def VarRNNStep():
#     def forward(input, hidden, cell, mask):
# def StackedStep(layer, num_layers, lstm=False):
#     def forward(input, hidden, cells, mask):
# def AutogradVarRNNStep(num_layers=1, lstm=False):
#     def forward(input, cells, hidden, mask):


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

__author__ = 'max'



class VarRNNBase(nn.Module):
    def __init__(self, Cell, input_size, hidden_size,
                 num_layers=1, bias=True, batch_first=False,
                 dropout=(0, 0), bidirectional=False, **kwargs):

        super(VarRNNBase, self).__init__()
        self.Cell = Cell
        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: