Repository: rwth-i6/returnn

Issue:
Functional RecLayer variant
Which does not have paramaters but all weights etc are inputs from other layers.


Discussion:
Maybe a more generic variant is also to use the `reuse_params` mechanism for this, and let this return other layer outputs.


Resolving patch:
diff --git a/returnn/tf/layers/base.py b/returnn/tf/layers/base.py
--- a/returnn/tf/layers/base.py
+++ b/returnn/tf/layers/base.py
@@ -1072,8 +1072,11 @@ def add_param(self, param, custom_update=None, trainable=None, saveable=None, ax
       custom_update.set_on_var(param)
     if axes_split_info:
       tf_util.set_param_axes_split_info(param, axes_split_info)
+    name_scope_prefix = self.get_absolute_name_scope_prefix()
     if self.reuse_params:
-      name_scope_prefix = self.reuse_params.get_absolute_name_scope_prefix(base_layer=self, param=param)
+      if not param.name.startswith(name_scope_prefix):
+        # We likely used an existing var from elsewhere. Return as-is.
+        return _param
     else:
       name_scope_prefix = self.get_absolute_name_scope_prefix()
     assert param.name
@@ -1770,8 +1773,10 @@ def optional_get_layer(layer_name):
         else:
           assert isinstance(value, dict)
           value = value.copy()
-          if "reuse_layer" in value:
+          if value.get("reuse_layer", None):
             value["reuse_layer"] = optional_get_layer(value["reuse_layer"])
+          if value.get("layer_output", None):
+            value["layer_output"] = get_layer(value["layer_output"])  # not optional, we need it right away
         opts["map"][key] = ReuseParams(**value)
     return ReuseParams(**opts)
 
@@ -1902,18 +1907,22 @@ def opt_get_layer(layer_name):
           name=layer_name, layer_class=layer_class, **layer_desc)
 
   # noinspection PyShadowingBuiltins
-  def __init__(self, reuse_layer=None, map=None, custom=None, auto_create_missing=False):
+  def __init__(self, reuse_layer=None, map=None, custom=None, auto_create_missing=False, layer_output=None, shape=None):
     """
     :param LayerBase|ReuseParams.LazyLayerResolver|None reuse_layer:
     :param dict[str,ReuseParams]|None map:
     :param (**kwargs)->(tf.Tensor|tf.Variable) custom: see :func:`self.variable_custom_getter`
     :param bool auto_create_missing:
+    :param LayerBase|None layer_output:
+    :param tuple[DimensionTag]|None shape:
     """
     assert isinstance(reuse_layer, (LayerBase, ReuseParams.LazyLayerResolver)) or not reuse_layer
     self._reuse_layer = reuse_layer
     self.param_map = map
     self.custom_func = custom
     self.auto_create_missing = auto_create_missing
+    self.layer_output = layer_output
+    self.shape = shape
 
   def __repr__(self):
     return "<%s reuse_layer %r, map %r>" % (self.__class__.__name__, self._reuse_layer, self.param_map)
@@ -1930,43 +1939,6 @@ def reuse_layer(self):
       return self._reuse_layer
     return None
 
-  def get_absolute_name_scope_prefix(self, base_layer, param):
-    """
-    :param LayerBase base_layer:
-    :param tf.Variable param: e.g. "base_layer/rec/W"
-    :return: e.g. "base_layer/" or "base_layer/rec/", always with "/" at end
-    :rtype: str
-    """
-    # abs_scope_prefix: e.g. "current_layer/" or "current_layer/rec/"
-    abs_scope_prefix = base_layer.get_absolute_name_scope_prefix()
-    assert abs_scope_prefix.endswith("/")
-    from returnn.tf.util.basic import get_current_var_scope_name
-    # abs_cur_scope: e.g. "current_layer/rec/" or "current_layer/" or "current_layer/batch_norm/"
-    abs_cur_scope = get_current_var_scope_name() + "/"
-    assert abs_cur_scope.startswith(abs_scope_prefix)
-    cur_scope_rel = abs_cur_scope[len(abs_scope_prefix):]  # e.g. "batch_norm/" or ""
-    assert cur_scope_rel == "" or abs_cur_scope.endswith("/")
-    assert param.name[-2:] == ":0"
-    abs_param_name = param.name[:-2]
-    abs_param_scope, _sep, param_name = abs_param_name.rpartition("/")
-    assert param_name and (abs_param_scope + _sep).endswith(cur_scope_rel)
-    if cur_scope_rel:
-      param_name = cur_scope_rel + param_name
-    del abs_param_scope
-    assert abs_param_name.endswith("/" + param_name)
-    if self.custom_func:  # Could be any base absolute name scope prefix, so just return what we have.
-      return abs_param_name[:-len(param_name)]
-    if self.param_map is not None and param_name in self.param_map:
-      return self.param_map[param_name].get_absolute_name_scope_prefix(base_layer=base_layer, param=param)
-    if self.reuse_layer and param_name in self.reuse_layer.params:
-      reuse_layer_prefix = self.reuse_layer.get_absolute_name_scope_prefix()
-      assert reuse_layer_prefix + param_name == abs_param_name
-      return reuse_layer_prefix
-    assert self.auto_create_missing
-    base_layer_prefix = base_layer.get_absolute_name_scope_prefix()
-    assert base_layer_prefix + param_name == abs_param_name
-    return base_layer_prefix
-
   def get_variable_scope(self, base_layer, **kwargs):
     """
     :param LayerBase base_layer:
@@ -1979,7 +1951,7 @@ def _variable_custom_getter(**kwargs_):
           tf_compat.v1.get_variable_scope(), custom_getter=_variable_custom_getter, **kwargs) as scope:
       return scope
 
-  def variable_custom_getter(self, getter, name, base_layer, **kwargs):
+  def variable_custom_getter(self, base_layer, name, shape, dtype, getter, **kwargs):
     """
     By TF docs, from :func:`_VariableStore.get_variable`:
     Callable that takes as a first argument the true getter,
@@ -1996,31 +1968,42 @@ def custom_getter(getter, name, *args, **kwargs):
     ```
     In addition, we get the argument `base_scope_name`, via :func:`self.get_variable_scope`.
 
-    :param (...)->tf.Variable getter:
-    :param str name: absolute param name
     :param LayerBase base_layer: we expect that this is the prefix of ``name``
+    :param str name: absolute param name
+    :param tuple[int] shape:
+    :param tensorflow.DType dtype:
+    :param (...)->tf.Variable getter:
     :rtype: tf.Variable|tf.Tensor
     """
+    if self.shape is not None:
+      assert shape == tuple(d.dimension for d in self.shape), "unexpected shape %r for param %r" % (shape, name)
     abs_scope_prefix = base_layer.get_absolute_name_scope_prefix()
     assert not abs_scope_prefix or abs_scope_prefix.endswith("/")
     assert name.startswith(abs_scope_prefix)
     param_name = name[len(abs_scope_prefix):]  # e.g. "W" (not "rec/W")
     if self.custom_func:
       return self.custom_func(
-        base_layer=base_layer, reuse_layer=self.reuse_layer, name=param_name, getter=getter, full_name=name, **kwargs)
+        base_layer=base_layer, reuse_layer=self.reuse_layer, full_name=name,
+        name=param_name, shape=shape, dtype=dtype, getter=getter, **kwargs)
     if self.param_map is not None:
       if not self.auto_create_missing:
         assert param_name in self.param_map
       if param_name in self.param_map:
         return self.param_map[param_name].variable_custom_getter(
-          getter=getter, name=name, base_layer=base_layer, **kwargs)
+          base_layer=base_layer, name=name, shape=shape, dtype=dtype, getter=getter, **kwargs)
     if self.reuse_layer:
       if not self.auto_create_missing:
         assert param_name in self.reuse_layer.params
       if param_name in self.reuse_layer.params:
         return self.reuse_layer.params[param_name]
+    if self.layer_output:
+      if self.shape is not None:
+        out = self.layer_output.output.copy_compatible_to(Data(name=name, dim_tags=self.shape, dtype=str(dtype)))
+        return out.placeholder
+      assert shape == self.layer_output.output.batch_shape
+      return self.layer_output.output.placeholder
     assert self.auto_create_missing
-    return getter(name=name, **kwargs)
+    return getter(name=name, shape=shape, dtype=dtype, **kwargs)
 
 
 class SearchChoices(object):


Test patch:
diff --git a/tests/test_TFEngine.py b/tests/test_TFEngine.py
--- a/tests/test_TFEngine.py
+++ b/tests/test_TFEngine.py
@@ -2973,16 +2973,12 @@ def test_preload_from_files_with_reuse():
     for param_name in ["W", "b"]:
       param_orig = layer_orig[param_name]
       param_clone_main = layer_clone_main[param_name]
-      param_clone_clone = layer_clone_clone[param_name]
-      numpy.testing.assert_array_equal(param_orig, param_clone_clone)
       numpy.testing.assert_array_equal(param_orig, param_clone_main)
 
     main = engine.network.layers["main_" + layer_name]
     clone = engine.network.layers["clone_" + layer_name]
     assert_equal(set(main.params.keys()), {"W", "b"})
-    assert_equal(set(clone.params.keys()), {"W", "b"})
-    assert main.params["W"] is clone.params["W"]
-    assert main.params["b"] is clone.params["b"]
+    assert_equal(set(clone.params.keys()), set())
 
   engine.finalize()
 
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -192,6 +192,46 @@ def test_LinearLayer_two_time_dims_allow_broadcast_all_sources():
     session.run(fetches=output.placeholder, feed_dict=make_feed_dict(network.extern_data))
 
 
+def test_LinearLayer_reuse_params_layer_output():
+  from returnn.tf.util.data import batch_dim
+  with make_scope() as session:
+    time_dim = SpatialDim("time")
+    data_feat_dim = FeatureDim("feature", dimension=5)
+    out_feat_dim = FeatureDim("feature", dimension=7)
+    config = Config({
+      "extern_data": {
+        "data": {"dim_tags": [batch_dim, time_dim, data_feat_dim]},
+      },
+    })
+    network = TFNetwork(config=config)
+    network.construct_from_dict({
+      "weights": {"class": "variable", "shape": [data_feat_dim, out_feat_dim]},
+      "bias": {"class": "variable", "shape": [out_feat_dim]},
+      "out1": {
+        "class": "linear", "from": "data", "out_dim": out_feat_dim,
+        "is_output_layer": True,
+        "reuse_params": {
+          "map": {
+            "W": {"layer_output": "weights"},
+            "b": {"layer_output": "bias"}}}},
+      "out2_": {
+        "class": "dot", "from": ["data", "weights"],
+        "red1": data_feat_dim, "red2": data_feat_dim,
+        "var1": [batch_dim, time_dim], "var2": out_feat_dim},
+      "out2": {
+        "class": "combine", "kind": "add", "from": ["out2_", "bias"],
+        "is_output_layer": True}
+    })
+    out1 = network.get_layer("out1").output
+    out2 = network.get_layer("out2").output
+    params = network.get_params_list()
+    assert len(params) == 2  # weights and bias
+    session.run(tf_compat.v1.global_variables_initializer())
+    out1_np, out2_np = session.run(
+      fetches=(out1.placeholder, out2.placeholder), feed_dict=make_feed_dict(network.extern_data))
+    numpy.testing.assert_array_equal(out1_np, out2_np)
+
+
 def test_PadLayer_time():
   n_batch, n_time, n_in = 7, 3, 20
   config = Config({
@@ -2851,9 +2891,7 @@ def test_reuse_params():
     l1 = network.layers["l1"]
     l2 = network.layers["output"]
     assert_equal(set(l1.params.keys()), {"W", "b"})
-    assert_equal(set(l2.params.keys()), {"W", "b"})
-    assert l1.params["W"] is l2.params["W"]
-    assert l1.params["b"] is l2.params["b"]
+    assert_equal(set(l2.params.keys()), set())
     assert_equal(set(network.get_trainable_params()), {l1.params["W"], l1.params["b"]})
 
 
@@ -4374,9 +4412,7 @@ def test_ReuseParams_dep_loop():
   l2 = network.get_layer("layer2")
   assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
   assert isinstance(l2, LinearLayer)
-  assert set(l1.params.keys()) == set(l2.params.keys()) == {"W", "b"}
-  assert l1.params["W"] is l2.params["W"]
-  assert l1.params["b"] is l2.params["b"]
+  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}
 
   def make_feed_dict(seq_len=10):
     random = numpy.random.RandomState(seed=1)
@@ -4439,9 +4475,7 @@ def test_ReuseParams_dep_loop_2():
   l2 = network.get_layer("layer2")
   assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
   assert isinstance(l2, LinearLayer)
-  assert set(l1.params.keys()) == set(l2.params.keys()) == {"W", "b"}
-  assert l1.params["W"] is l2.params["W"]
-  assert l1.params["b"] is l2.params["b"]
+  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}
 
   def make_feed_dict(seq_len=10):
     random = numpy.random.RandomState(seed=1)
@@ -4504,9 +4538,7 @@ def test_ReuseParams_dep_loop_3():
   l2 = network.get_layer("layer2")
   assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
   assert isinstance(l2, LinearLayer)
-  assert set(l1.params.keys()) == set(l2.params.keys()) == {"W", "b"}
-  assert l1.params["W"] is l2.params["W"]
-  assert l1.params["b"] is l2.params["b"]
+  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}
 
   def make_feed_dict(seq_len=10):
     random = numpy.random.RandomState(seed=1)