cleanup

No rec ...Step anymore.
Fix #81.

Author: albertz

nn/rec.py

@@ -2,7 +2,6 @@
 Basic RNNs.
 """
 
-from typing import Optional
 from .. import nn
 from ._generated_layers import _Rec
 
@@ -14,52 +13,13 @@ class LSTM(_Rec):
   def __init__(self, out_dim: nn.Dim, **kwargs):
     super().__init__(out_dim=out_dim, unit="nativelstm2", **kwargs)
 
-  def __call__(
-        self, source: nn.LayerRef, *, axis: nn.Dim, initial_state: Optional[nn.LayerState] = None) -> nn.Layer:
-    """make layer"""
-    return super()(source, axis=axis, initial_state=initial_state)
-
-
-class LSTMStep(_Rec):
-  """
-  LSTM operating one step. returns (output, state) tuple, where state is (h,c).
-  """
-  default_name = "lstm"  # make consistent to LSTM
-
-  def __init__(self, out_dim: nn.Dim, **kwargs):
-    super().__init__(out_dim=out_dim, unit="nativelstm2", **kwargs)
-
-  def __call__(self, source: nn.LayerRef, *, state: nn.LayerState) -> nn.Layer:
-    """make layer"""
-    return super()(source, state=state, axis=nn.single_step_dim)
-
 
 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):
+  def __init__(self, n_out: int, zoneout_factor_cell: float = 0., zoneout_factor_output: float = 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)
-
-  def __call__(
-        self, source: nn.LayerRef, *, axis: nn.Dim, initial_state: Optional[nn.LayerState] = None) -> nn.Layer:
-    """make layer"""
-    return super()(source, axis=axis, 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)
-
-  def __call__(self, source: nn.LayerRef, *, state: nn.LayerState) -> nn.Layer:
-    """make layer"""
-    return super()(source, state=state, axis=nn.single_step_dim)
+      unit_opts={'zoneout_factor_cell': zoneout_factor_cell, 'zoneout_factor_output': zoneout_factor_output},
+      **kwargs)

tests/test_models_rec.py

@@ -30,9 +30,9 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
       """
       # https://github.com/rwth-i6/returnn_common/issues/16
       with nn.Loop() as loop:
-        x_ = loop.unstack(x, axis="T", declare_rec_time=True)
+        x_ = loop.unstack(x, axis="T", declare_rec_time=True)  # TODO how to get axis?
         loop.state.h = nn.State(initial=0)  # TODO proper initial...
-        loop.state.h = self.rec_linear(nn.concat((x_, "F"), (loop.state.h, "F")))
+        loop.state.h = self.rec_linear(nn.concat((x_, "F"), (loop.state.h, self.rec_linear.out_dim)))  # TODO dim?
         y = loop.stack(loop.state.h)
       return y
 
@@ -59,7 +59,9 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
 
 def test_lstm_default_name():
   assert_equal(nn.LSTM(nn.FeatureDim("out", 3)).get_default_name(), "lstm")
-  assert_equal(nn.LSTMStep(nn.FeatureDim("out", 3)).get_default_name(), "lstm")
+  # no LSTMStep anymore, so nothing really to test here.
+  # https://github.com/rwth-i6/returnn_common/issues/81
+  # assert_equal(nn.LSTMStep(nn.FeatureDim("out", 3)).get_default_name(), "lstm")
 
 
 def test_rec_inner_lstm():
@@ -74,7 +76,7 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
       Forward
       """
       with nn.Loop() as loop:
-        x_ = loop.unstack(x, axis="T", declare_rec_time=True)
+        x_ = loop.unstack(x, axis="T", declare_rec_time=True)  # TODO how to get axis?
         loop.state.lstm = nn.State(initial=self.lstm.default_initial_state())
         y_, loop.state.lstm = self.lstm(x_, state=loop.state.lstm)
         y = loop.stack(y_)
@@ -98,7 +100,7 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
         loop.state.i = nn.State(initial=0.)
         loop.state.i = loop.state.i + 1.
         loop.end(loop.state.i >= 5., include_eos=True)
-        y = loop.stack(loop.state.i * nn.reduce(x, mode="mean", axis="T"))
+        y = loop.stack(loop.state.i * nn.reduce(x, mode="mean", axis="T"))  # TODO axis
       return y
 
   net = _Net()
@@ -118,9 +120,9 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
       """
       Forward
       """
-      y, state = self.lstm(x)
-      y_ = nn.reduce(y, mode="mean", axis="T")  # TODO just because concat allow_broadcast=True does not work yet...
-      res = nn.concat((y_, "F"), (state.h, "F"), (state.c, "F"))
+      y, state = self.lstm(x)  # TODO axis
+      res = nn.concat(
+        (y, self.lstm.out_dim), (state.h, self.lstm.out_dim), (state.c, self.lstm.out_dim), allow_broadcast=True)
       return res
 
   net = _Net()
@@ -144,7 +146,7 @@ def __call__(self, x: nn.LayerRef) -> nn.LayerRef:
       y = self.linear(x)
       state = None
       for _ in range_(3):
-        y, state = self.lstm(y, initial_state=state)
+        y, state = self.lstm(y, initial_state=state)  # TODO axis?
       return y
 
   net = _Net()