Repository: rwth-i6/returnn

Issue:
GatherLayer: problem with search beam
There is a problem with the `get_out_data_from_opts` function of `GatherLayer` during search. This is the problematic part:

```
out_type = input_data.get_kwargs(include_special_axes=False)
out_type["name"] = "%s_output" % name
out_type["dim_tags"] = dim_tags
out_type["beam"] = SearchBeam.get_combined_beam(input_data.beam, position_data.beam)
out_type["available_for_inference"] = input_data.available_for_inference and position_data.available_for_inference
...
output_data = Data(**out_type)
```

`out_type` uses the `batch` information from `input_data` and the **combined** `beam` information from `input_data` and `position_data`. This means that `batch.beam` and `beam` can be different if `input_data` and `position_data` have different beam information and this leads to a problem inside of `output_data = Data(**out_type)`:

```
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 1490, in Data.__init__
    line: assert batch.beam == beam
    locals:
      batch = <local> BatchInfo{B}
      batch.beam = <local> None
      beam = <local> SearchBeam(name='output/prev:output', beam_size=3)
AssertionError
```

The full error and the code to reproduce the problem are in https://gist.github.com/robin-p-schmitt/9e8d968fa2609cf17909c1bd6e25c5c0.


Discussion:
See for example `DotLayer`. When you add sth like:
```
out_type["batch"] = BatchInfo.get_common_batch_info([src.batch for src in (input_data, position_data)])
```
This probably fixes the problem.
Can you try, and maybe create a PR? The PR should then also include your test.


Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -1155,6 +1155,7 @@ def get_out_data_from_opts(cls, name, sources, position, axis, **kwargs):
     :param str axis:
     :rtype: Data
     """
+    from returnn.tf.util.data import BatchInfo
     input_data = get_concat_sources_data_template(sources)
     old_gather_axis = input_data.get_axis_from_description(axis, allow_int=False)
 
@@ -1183,6 +1184,7 @@ def get_out_data_from_opts(cls, name, sources, position, axis, **kwargs):
     out_type["dim_tags"] = dim_tags
     out_type["beam"] = SearchBeam.get_combined_beam(input_data.beam, position_data.beam)
     out_type["available_for_inference"] = input_data.available_for_inference and position_data.available_for_inference
+    out_type["batch"] = BatchInfo.get_common_batch_info([src.batch for src in (input_data, position_data)])
 
     # Take axes from input_data if they exist there, otherwise from position_data
     for axis_kind in Data.SpecialAxesNames:


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -2792,6 +2792,29 @@ def test_GatherLayer_constant_position():
           np.testing.assert_almost_equal(out_seqs[b, f1, f2], values_seqs[b, f1, position, f2])
 
 
+def test_GatherLayer_search_beam():
+  from returnn.tf.network import TFNetwork
+  from returnn.config import Config
+  with make_scope() as session:
+    n_out = 5
+    config = Config({
+      "debug_print_layer_output_template": True,
+      "extern_data": {
+        "data": {"dim": n_out},
+        "classes": {"dim": n_out, "sparse": True}
+      }})
+    net = TFNetwork(config=config, search_flag=True)
+    net.construct_from_dict({
+      "output": {
+        "class": "rec", "from": "data:data", "unit": {
+          "position": {"class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
+          "gather": {"class": "gather", "from": "base:data:data", "position": "position", "axis": "t"},  # [B,T,slice,D]
+          "prob": {"class": "softmax", "from": "gather", "target": "classes", "loss": "ce"},
+          'output': {
+            'class': 'choice', 'target': "classes", 'beam_size': 3, 'from': "prob", "input_type": "prob",
+            "initial_output": 0}}}})
+
+
 def test_SliceNdLayer():
   n_batch = 5
   n_time = 7