returnn 1.20230703.180607__tar.gz → 1.20230703.193930__tar.gz

{returnn-1.20230703.180607/returnn.egg-info → returnn-1.20230703.193930}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230703.180607
+Version: 1.20230703.193930
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230703.193930/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230703.193930'
+long_version = '1.20230703.193930+git.3d14839'

{returnn-1.20230703.180607 → returnn-1.20230703.193930}/returnn/tf/layers/rec.py

@@ -4,6 +4,7 @@ Defines multiple recurrent layers, most importantly :class:`RecLayer`.
 
 from __future__ import annotations
 
+import contextlib
 import typing
 import tensorflow as tf
 import returnn.tf.compat as tf_compat
@@ -5450,15 +5451,16 @@ class RnnCellLayer(_ConcatInputLayer):
 
     @classmethod
     def get_rec_initial_state_inner(
-        cls, initial_shape, name, state_key="state", key=None, initial_state=None, shape_invariant=None, rec_layer=None
+        cls, initial_shape, name, state_key=None, key=None, initial_state=None, shape_invariant=None, rec_layer=None
     ):
         """
         Generate initial hidden state. Primarily used as a inner function for RnnCellLayer.get_rec_initial_state.
 
         :param tuple initial_shape: shape of the initial state.
         :param str name: layer name.
-        :param str state_key: key to be used to get the state from final_rec_vars.
-        :param str|None key: key/attribute of the state if state is a dictionary/namedtuple
+        :param str|None state_key: key to be used to get the state from final_rec_vars.
+            "state" by default.
+        :param str|int|None key: key/attribute of the state if state is a dictionary/namedtuple
             (like 'c' and 'h' for LSTM states).
         :param LayerBase|str|int|float|None|list|tuple|namedtuple initial_state: see code
         :param tuple shape_invariant: If provided, directly used.
@@ -5473,13 +5475,14 @@ class RnnCellLayer(_ConcatInputLayer):
             # https://github.com/rwth-i6/returnn/pull/1248#issuecomment-1618125397
             if isinstance(rec_layer, RecLayer) and isinstance(rec_layer.cell, _SubnetworkRecCell):
                 final_rec_vars = rec_layer.cell.get_final_rec_vars(name)
-                last_state = cls.get_state_by_key(final_rec_vars[state_key], key=key, shape=initial_shape)
+                last_state = cls.get_state_by_key(final_rec_vars[state_key or "state"], key=key, shape=initial_shape)
             else:
-                last_state = rec_layer.get_last_hidden_state(key=key)
+                assert key is None or not state_key
+                last_state = rec_layer.get_last_hidden_state(key=key if key is not None else state_key)
             last_state.set_shape(shape_invariant)
             return last_state
 
-        if state_key == "state":
+        if not state_key:
             if key is None:
                 key_name = "var"
             else:
@@ -5489,7 +5492,6 @@ class RnnCellLayer(_ConcatInputLayer):
                 key_name = state_key
             else:
                 key_name = f"{state_key}_{key}"
-        from returnn.util.basic import dummy_noop_ctx
 
         if shape_invariant is None:
             shape_invariant = tuple([d if isinstance(d, int) and d != 0 else None for d in initial_shape])
@@ -5510,7 +5512,7 @@ class RnnCellLayer(_ConcatInputLayer):
         elif initial_state == "var":  # Initial state is a trainable variable.
             # Assume the first dimension to be batch_dim.
             assert shape_invariant[0] is None and all([d is not None for d in shape_invariant[1:]])
-            with rec_layer.var_creation_scope() if rec_layer else dummy_noop_ctx():
+            with rec_layer.var_creation_scope() if rec_layer else contextlib.nullcontext():
                 var = tf_compat.v1.get_variable(
                     "initial_%s" % key_name, shape=initial_shape[1:], initializer=tf.zeros_initializer()
                 )

{returnn-1.20230703.180607 → returnn-1.20230703.193930}/returnn/util/basic.py

@@ -4186,14 +4186,6 @@ def ensure_list_of_type(ls, type_):
     return ls
 
 
-@contextlib.contextmanager
-def dummy_noop_ctx():
-    """
-    Provides a no-op context manager.
-    """
-    yield None
-
-
 def _get_ngrams(segment, max_order):
     """Extracts all n-grams upto a given maximum order from an input segment.
     Code adapted from Google Tensor2Tensor.

{returnn-1.20230703.180607 → returnn-1.20230703.193930/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230703.180607
+Version: 1.20230703.193930
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20230703.180607 → returnn-1.20230703.193930}/tests/test_TFNetworkRecLayer.py

@@ -11455,61 +11455,77 @@ def test_self_attention_keep_over_epoch_state():
 
 
 def test_self_attention_placeholder_state():
+    n_batch = 3
     num_heads = 2
     total_key_dim = 6
     total_value_dim = 14
-    with make_scope() as session:
-        n_in, n_out = 2, 3
-        config = Config(
-            {
-                "extern_data": {"data": {"dim": n_in}, "classes": {"dim": n_out, "sparse": True}},
-            }
-        )
-        net_dict = {
-            "loop": {
-                "class": "rec",
-                "from": "data",
-                "unit": {
-                    "self_att": {
-                        "class": "self_attention",
-                        "from": "data:source",
-                        "attention_left_only": True,
-                        "num_heads": num_heads,
-                        "total_key_dim": total_key_dim,
-                        "n_out": total_value_dim,
-                        "initial_state": "placeholder",
+    rnd = numpy.random.RandomState(42)
+    prev_time = 2
+    k_init_v = rnd.normal(size=(n_batch, num_heads, prev_time, total_key_dim // num_heads)).astype("float32")
+    v_init_v = rnd.normal(size=(n_batch, num_heads, prev_time, total_value_dim // num_heads)).astype("float32")
+    prev_params_serialized = None
+    prev_out_vs = None
+    for opt in [False, True]:
+        with make_scope() as session:
+            n_in, n_out = 2, 3
+            config = Config(
+                {
+                    "extern_data": {"data": {"dim": n_in}, "classes": {"dim": n_out, "sparse": True}},
+                }
+            )
+            net_dict = {
+                "loop": {
+                    "class": "rec",
+                    "from": "data",
+                    "unit": {
+                        "self_att": {
+                            "class": "self_attention",
+                            "from": "data:source",
+                            "attention_left_only": True,
+                            "num_heads": num_heads,
+                            "total_key_dim": total_key_dim,
+                            "n_out": total_value_dim,
+                            "initial_state": "placeholder",
+                        },
+                        "output": {"class": "copy", "from": "self_att"},
                     },
-                    "output": {"class": "copy", "from": "self_att"},
+                    "optimize_move_layers_out": opt,
                 },
-                "optimize_move_layers_out": False,
-            },
-            "output": {"class": "softmax", "loss": "ce", "target": "classes", "from": "loop", "n_out": n_out},
-        }
-        network = TFNetwork(config=config)
-        network.construct_from_dict(net_dict)
+                "output": {"class": "softmax", "loss": "ce", "target": "classes", "from": "loop", "n_out": n_out},
+            }
+            network = TFNetwork(config=config)
+            network.construct_from_dict(net_dict)
 
-        expected_model_params = {"loop/rec/self_att/QKV", "output/W", "output/b"}
+            expected_model_params = {"loop/rec/self_att/QKV", "output/W", "output/b"}
 
-        print("All global variables:")
-        params = tf_compat.v1.global_variables()
-        params = {v.op.name: v for v in params}
-        pprint(params)
-        params.pop("global_step", None)
-        assert set(params.keys()) == expected_model_params
+            print("All global variables:")
+            params = tf_compat.v1.global_variables()
+            params = {v.op.name: v for v in params}
+            pprint(params)
+            params.pop("global_step", None)
+            assert set(params.keys()) == expected_model_params
 
-        from test_TFNetworkLayer import make_feed_dict
+            if prev_params_serialized is None:
+                network.initialize_params(session)
+                prev_params_serialized = network.get_params_serialized(session)
+            else:
+                network.set_params_by_serialized(prev_params_serialized, session)
 
-        network.initialize_params(session)
-        out = network.get_default_output_layer().output
-        n_batch = 3
-        feed_dict = make_feed_dict(network.extern_data, n_batch=n_batch)
-        k_init = session.graph.get_tensor_by_name("loop/rec/self_att/initial_state_placeholder_k_left:0")
-        v_init = session.graph.get_tensor_by_name("loop/rec/self_att/initial_state_placeholder_v_left:0")
-        k_last = session.graph.get_tensor_by_name("loop/rec/self_att/last_state_k_left:0")
-        v_last = session.graph.get_tensor_by_name("loop/rec/self_att/last_state_v_left:0")
-        feed_dict[k_init] = numpy.zeros((n_batch, num_heads, 0, total_key_dim // num_heads), dtype="float32")
-        feed_dict[v_init] = numpy.zeros((n_batch, num_heads, 0, total_value_dim // num_heads), dtype="float32")
-        session.run((out.placeholder, k_last, v_last), feed_dict=feed_dict)
+            from test_TFNetworkLayer import make_feed_dict
+
+            out = network.get_default_output_layer().output
+            feed_dict = make_feed_dict(network.extern_data, n_batch=n_batch)
+            k_init = session.graph.get_tensor_by_name("loop/rec/self_att/initial_state_placeholder_k_left:0")
+            v_init = session.graph.get_tensor_by_name("loop/rec/self_att/initial_state_placeholder_v_left:0")
+            k_last = session.graph.get_tensor_by_name("loop/rec/self_att/last_state_k_left:0")
+            v_last = session.graph.get_tensor_by_name("loop/rec/self_att/last_state_v_left:0")
+            feed_dict[k_init] = k_init_v
+            feed_dict[v_init] = v_init_v
+            out_vs = session.run((out.placeholder, k_last, v_last), feed_dict=feed_dict)
+            if prev_out_vs is not None:
+                for v, prev_v in zip(out_vs, prev_out_vs):
+                    assert_allclose(v, prev_v, atol=1e-5, rtol=1e-5)
+            prev_out_vs = out_vs
 
 
 def test_OptimalCompletionsLayer():

returnn-1.20230703.180607/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230703.180607'
-long_version = '1.20230703.180607+git.5f239b4'