# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn as nn


class LSTMCellWithZoneOut(nn.Module):
    """
    Zoneout: Regularizing RNNs by Randomly Preserving Hidden Activations
    https://arxiv.org/abs/1606.01305
    """

    def __init__(self, prob: float, input_size: int, hidden_size: int,
                 bias: bool = True):
        super(LSTMCellWithZoneOut, self).__init__()
        self.lstm_cell = nn.LSTMCell(input_size, hidden_size, bias=bias)
        self.prob = prob
        if prob > 1.0 or prob < 0.0:
            raise ValueError("zoneout probability must be in the range from "
                             "0.0 to 1.0.")

    def zoneout(self, h, next_h, prob):
        if isinstance(h, tuple):
            return tuple(
                [self.zoneout(h[i], next_h[i], prob) for i in range(len(h))]
            )

        if self.training:
            mask = h.new_zeros(*h.size()).bernoulli_(prob)
            return mask * h + (1 - mask) * next_h

        return prob * h + (1 - prob) * next_h

    def forward(self, x, h):
        return self.zoneout(h, self.lstm_cell(x, h), self.prob)