# 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
        self.noise_hidden = None

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weight_hh)
        nn.init.xavier_uniform_(self.weight_ih)
        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)
                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):
        if self.nonlinearity == "tanh":
            func = rnn_F.VarRNNTanhCell
        elif self.nonlinearity == "relu":

# 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

        self.noise_hidden = None

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weight_hh)
        nn.init.xavier_uniform_(self.weight_ih)
        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)
                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):
        if self.nonlinearity == "tanh":
            func = rnn_F.VarRNNTanhCell
        elif self.nonlinearity == "relu":

            func = rnn_F.VarRNNReLUCell