added zoneout to rec.py (#86)

Author: Atticus1806

nn/rec.py

@@ -37,3 +37,37 @@ def make_layer_dict(
     # TODO specify per-step, how? this should also work without rec loop, when there is no time dim.
     #  https://github.com/rwth-i6/returnn/issues/847
     return super().make_layer_dict(source, state=state, axis=None)
+
+
+class ZoneoutLSTM(_Rec):
+  """
+  LSTM with zoneout operating on a sequence. returns (output, final_state) tuple, where final_state is (h,c).
+  """
+  def __init__(self, n_out: int, zoneout_factor_cell: int = 0., zoneout_factor_output: int = 0., **kwargs):
+    super().__init__(
+      n_out=n_out, unit="zoneoutlstm",
+      unit_opts={'zoneout_factor_cell': zoneout_factor_cell, 'zoneout_factor_output': zoneout_factor_output}, **kwargs)
+
+  # noinspection PyMethodOverriding
+  def make_layer_dict(
+      self, source: nn.LayerRef, *, initial_state: Optional[nn.LayerState] = None) -> nn.LayerDictRaw:
+    """make layer"""
+    return super().make_layer_dict(source, initial_state=initial_state)
+
+
+class ZoneoutLSTMStep(_Rec):
+  """
+  LSTM with zoneout operating one step. returns (output, state) tuple, where state is (h,c).
+  """
+  default_name = "zoneoutlstm"  # make consistent to ZoneoutLSTM
+
+  def __init__(self, n_out: int, zoneout_factor_cell: int = 0., zoneout_factor_output: int = 0., **kwargs):
+    super().__init__(
+      n_out=n_out, unit="zoneoutlstm",
+      unit_opts={'zoneout_factor_cell': zoneout_factor_cell, 'zoneout_factor_output': zoneout_factor_output}, **kwargs)
+
+  # noinspection PyMethodOverriding
+  def make_layer_dict(
+        self, source: nn.LayerRef, *, state: nn.LayerState) -> nn.LayerDictRaw:
+    """make layer"""
+    return super().make_layer_dict(source, state=state)