added zoneout to rec

Author: Atticus1806

nn/rec.py

@@ -35,3 +35,33 @@ def make_layer_dict(
         self, source: nn.LayerRef, *, state: nn.LayerState) -> nn.LayerDictRaw:
     """make layer"""
     return super().make_layer_dict(source, state=state)
+
+
+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, **kwargs):
+    super().__init__(n_out=n_out, unit="zoneoutlstm", **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, **kwargs):
+    super().__init__(n_out=n_out, unit="zoneoutlstm", **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)