single step

rwth-i6/returnn#847
#81

nn/base.py

@@ -48,9 +48,8 @@ def forward(self, x):
 
 from __future__ import annotations
 from typing import Dict, Any, Optional, List, Tuple, Union, Set, Iterator, Iterable
+from returnn.tf.util.data import *  # Dim, Data, and others
 from returnn.util.basic import NotSpecified, OptionalNotImplementedError
-# noinspection PyUnresolvedReferences
-from returnn.tf.util.data import Dim, SpatialDim, FeatureDim
 from tensorflow.python.util import nest
 
 

nn/rec.py

@@ -34,9 +34,7 @@ def __init__(self, out_dim: nn.Dim, **kwargs):
   def make_layer_dict(
         self, source: nn.LayerRef, *, state: nn.LayerState) -> nn.LayerDictRaw:
     """make layer"""
-    # 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)
+    return super().make_layer_dict(source, state=state, axis=nn.single_step_dim)
 
 
 class ZoneoutLSTM(_Rec):