LSTM with recurrent dropout.

  Args:
    hidden_size: the LSTM hidden size.
    keep_prob_c: the probability to use the new value of the cell state rather
      than freezing it.
    keep_prob_h: the probability to use the new value of the hidden state
      rather than freezing it.
    **kwargs: Extra keyword arguments to pass to the LSTM.

  Returns:
    A tuple (train_lstm, test_lstm) where train_lstm is an LSTM with
    recurrent dropout enabled to be used for training and test_lstm
    is the same LSTM without zoneout.

def lstm_with_zoneout(hidden_size, keep_prob_c=0.5, keep_prob_h=0.95, **kwargs):
  """LSTM with recurrent dropout.

  Args:
    hidden_size: the LSTM hidden size.
    keep_prob_c: the probability to use the new value of the cell state rather
      than freezing it.
    keep_prob_h: the probability to use the new value of the hidden state
      rather than freezing it.
    **kwargs: Extra keyword arguments to pass to the LSTM.

  Returns:
    A tuple (train_lstm, test_lstm) where train_lstm is an LSTM with
    recurrent dropout enabled to be used for training and test_lstm
    is the same LSTM without zoneout.
  """

  lstm = LSTM(hidden_size, **kwargs)
  keep_probs = LSTMState(keep_prob_h, keep_prob_c)
  train_lstm = ZoneoutWrapper(lstm, keep_probs, is_training=True)
  test_lstm = ZoneoutWrapper(lstm, keep_probs, is_training=False)
  return train_lstm, test_lstm