Repository: rwth-i6/returnn

Issue:
Problem with compile_tf_graph tool
Using the network dict shown in https://gist.github.com/robin-p-schmitt/3b3643efc7c8cfa8859f6d1a336549f9, with the following relevant layers:

```
"readout_in": {
                "activation": None,
                "class": "linear",
                "from": ["lm", "att"],
                "n_out": 1000,
            },
"lm": {
                "class": "rec",
                "from": ["prev:target_embed", "prev:att"],
                "n_out": 1024,
                "unit": "nativelstm2",
            },
"att": {"axes": "except_batch", "class": "merge_dims", "from": ["att0"]},
"att0": {
                "base": "base:enc_value",
                "class": "generic_attention",
                "weights": "att_weights",
            },
```

I get the following error:

```
...
layer /output(rec-subnet)/'readout_in': [B,F|F'readout_in:feature-dense'(1000)] float32
Mismatching shape: Tensor <tf.Tensor 'output/rec/concat_lm_att/concat_sources/concat:0' shape=(?, ?) dtype=float32> vs Data Data{'concat_lm_att', [B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}
fetch: <tf.Operation 'output/rec/concat_lm_att/concat_sources/concat' type=ConcatV2> [0], shape [None, None]
  inputs[0] 'values': <tf.Operation 'output/rec/lm/rec/concat_prev__target_embed_prev__att_squeeze_axes' type=Squeeze> [0], shape [None, 1024]
    inputs[0] 'input': <tf.Operation 'output/rec/lm/rec/NativeLstm2' type=NativeLstm2> [0], shape [1, None, 1024]
      inputs[0] 'x': <tf.Operation 'output/rec/lm/rec/add' type=AddV2> [0], shape [1, None, 4096]
        inputs[0] 'x': <tf.Operation 'output/rec/lm/rec/dot/Reshape_1' type=Reshape> [0], shape [1, None, 4096]
        inputs[1] 'y': <tf.Operation 'output/rec/lm/rec/add/ReadVariableOp' type=ReadVariableOp> [0], shape [4096]
      inputs[1] 'w': <tf.Operation 'output/rec/lm/rec/NativeLstm2/ReadVariableOp' type=ReadVariableOp> [0], shape [1024, 4096]
        inputs[0] 'resource': <tf.Operation 'output/rec/lm/rec/W_re' type=VarHandleOp> [0], shape []
      inputs[2] 'y0': <tf.Operation 'output/rec/while_loop_body/prev_extra/identity_lm_state_1' type=Identity> [0], shape [None, 1024]
        inputs[0] 'input': <tf.Operation 'output/rec/Identity_6' type=Identity> [0], shape [None, 1024]
      inputs[3] 'c0': <tf.Operation 'output/rec/while_loop_body/prev_extra/identity_lm_state_0' type=Identity> [0], shape [None, 1024]
        inputs[0] 'input': <tf.Operation 'output/rec/Identity_5' type=Identity> [0], shape [None, 1024]
      inputs[4] 'i': <tf.Operation 'output/rec/lm/rec/fixed_seq_len/expand_dims_unbroadcast/sequence_mask_time_major/cast_float32' type=Cast> [0], shape [None, None]
        inputs[0] 'x': <tf.Operation 'output/rec/lm/rec/fixed_seq_len/expand_dims_unbroadcast/sequence_mask_time_major/transpose' type=Transpose> [0], shape [None, None]
      inputs[5] 'start': <tf.Operation 'output/rec/lm/rec/start' type=Const value=0> [0], shape []
      inputs[6] 'step': <tf.Operation 'output/rec/lm/rec/step' type=Const value=1> [0], shape []
  inputs[1] 'axis': <tf.Operation 'output/rec/att/Reshape' type=Reshape> [0], shape [None, None]
    inputs[0] 'tensor': <tf.Operation 'output/rec/att/transpose' type=Transpose> [0], shape [None, 1, None, 2048]
      inputs[0] 'x': <tf.Operation 'output/rec/att0/Reshape_2' type=Reshape> [0], shape [None, 1, None, 2048]
        inputs[0] 'tensor': <tf.Operation 'output/rec/att0/MatMul' type=BatchMatMulV2> [0], shape [None, 1, None, 2048]
        inputs[1] 'shape': <tf.Operation 'output/rec/att0/Reshape_2/shape' type=Pack> [0], shape [4]
      inputs[1] 'perm': <tf.Operation 'output/rec/att/transpose/perm' type=Const value=[0 1 2 3]> [0], shape [4]
    inputs[1] 'shape': <tf.Operation 'output/rec/att/Reshape/shape' type=Pack> [0], shape [2]
      inputs[0] 'values': <tf.Operation 'output/rec/att/get_shape/strided_slice' type=StridedSlice> [0], shape []
        inputs[0] 'input': <tf.Operation 'output/rec/att/get_shape/Shape' type=Shape> [0], shape [4]
        inputs[1] 'begin': <tf.Operation 'output/rec/att/get_shape/strided_slice/stack' type=Const value=[0]> [0], shape [1]
        inputs[2] 'end': <tf.Operation 'output/rec/att/get_shape/strided_slice/stack_1' type=Const value=[1]> [0], shape [1]
        inputs[3] 'strides': <tf.Operation 'output/rec/att/get_shape/strided_slice/stack_2' type=Const value=[1]> [0], shape [1]
      inputs[1] '?0': <tf.Operation 'output/rec/att/Prod' type=Prod> [0], shape []
        inputs[0] 'input': <tf.Operation 'output/rec/att/Prod/input' type=Pack> [0], shape [3]
        inputs[1] 'reduction_indices': <tf.Operation 'output/rec/att/range' type=Range> [0], shape [1]
  inputs[2] '?0': <tf.Operation 'output/rec/concat_lm_att/concat_sources/concat/axis' type=Const value=1> [0], shape []
Exception creating layer /output(rec-subnet)/'readout_in' of class LinearLayer with opts:
{'_name': 'readout_in',
 '_network': <TFNetwork '/output(rec-subnet)' parent_layer=<RecStepByStepLayer 'output' out_type=Data{[T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}}> train=False>,
 'activation': None,
 'n_out': 1000,
 'name': 'readout_in',
 'network': <TFNetwork '/output(rec-subnet)' parent_layer=<RecStepByStepLayer 'output' out_type=Data{[T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}}> train=False>,
 'output': Data{'readout_in_output', [B,F|F'readout_in:feature-dense'(1000)], ctx=loop('dyn-time:output'[?])},
 'sources': [<RecLayer output/'lm' out_type=Data{[B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}>,
             <MergeDimsLayer output/'att' out_type=Data{[B,F|F'((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_fw:feature)+((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_bw:feature)'[B]{ctx=loop('dyn-time:output'[?])}], ctx=loop('dyn-time:output'[?])}>]}

ERROR: Got exception during in-loop construction of layer 'readout_in':
AssertionError: 

Template network (check out types / shapes):
output: <_TemplateLayer(ChoiceStateVarLayer)(:template:choice_state_var) output/'output' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])} (construction stack None)>
output_prob: <_TemplateLayer(SoftmaxLayer)(:template:softmax) output/'output_prob' out_type=Data{[B,F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])} (construction stack 'output')>
readout: <_TemplateLayer(ReduceOutLayer)(:template:reduce_out) output/'readout' out_type=Data{[B,F|F'(readout_in:feature-dense)//2'(500)], ctx=loop('dyn-time:output'[?])} (construction stack 'output_prob')>
readout_in: <_TemplateLayer(LinearLayer)(:template:linear) output/'readout_in' out_type=Data{[B,F|F'readout_in:feature-dense'(1000)], ctx=loop('dyn-time:output'[?])} (construction stack 'readout')>
lm: <_TemplateLayer(RecLayer)(:template:rec) output/'lm' out_type=Data{[B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'readout_in')>
target_embed: <_TemplateLayer(LinearLayer)(:template:linear) output/'target_embed' out_type=Data{[B,F|F'target_embed:feature-dense'(621)], ctx=loop('dyn-time:output'[?])} (construction stack 'lm')>
att: <_TemplateLayer(MergeDimsLayer)(:template:merge_dims) output/'att' out_type=Data{[B,F|F'((energy:feature-dense)*lstm5_fw:feature)+((energy:feature-dense)*lstm5_bw:feature)'(2048)], ctx=loop('dyn-time:output'[?])} (construction stack 'lm')>
att0: <_TemplateLayer(GenericAttentionLayer)(:template:generic_attention) output/'att0' out_type=Data{[B,F'energy:feature-dense'(1),F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)], ctx=loop('dyn-time:output'[?])} (construction stack 'att')>
att_weights: <_TemplateLayer(DropoutLayer)(:template:dropout) output/'att_weights' out_type=Data{[B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B]], ctx=loop('dyn-time:output'[?])} (construction stack 'att0')>
att_weights0: <_TemplateLayer(SoftmaxOverSpatialLayer)(:template:softmax_over_spatial) output/'att_weights0' out_type=Data{[B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B]], ctx=loop('dyn-time:output'[?])} (construction stack 'att_weights')>
energy: <_TemplateLayer(LinearLayer)(:template:linear) output/'energy' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1)], ctx=loop('dyn-time:output'[?])} (construction stack 'att_weights0')>
energy_tanh: <_TemplateLayer(ActivationLayer)(:template:activation) output/'energy_tanh' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy')>
energy_in: <_TemplateLayer(CombineLayer)(:template:combine) output/'energy_in' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy_tanh')>
weight_feedback: <_TemplateLayer(LinearLayer)(:template:linear) output/'weight_feedback' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'weight_feedback:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy_in')>
accum_att_weights: <_TemplateLayer(EvalLayer)(:template:eval) output/'accum_att_weights' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'inv_fertility:feature-dense'(1)], ctx=loop('dyn-time:output'[?])} (construction stack 'weight_feedback')>
att_query: <_TemplateLayer(LinearLayer)(:template:linear) output/'att_query' out_type=Data{[B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy_in')>
data:bpe: <_TemplateLayer(SourceLayer)(:template:source) output/'data:bpe' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, available_for_inference=False, ctx=loop('dyn-time:output'[?])} (construction stack 'output_prob')>
end: <_TemplateLayer(CompareLayer)(:template:compare) output/'end' out_type=Data{[B], dtype='bool', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])} (construction stack None)>

Exception creating layer /'output' of class RecStepByStepLayer with opts:
{'_name': 'output',
 '_network': <TFNetwork '' train=False>,
 '_orig_sources': [],
 '_target_layers': {'bpe': <DataNotAvailableLayer 'data:bpe' out_type=Data{[B,T|'time:var:extern_data:bpe'[B]], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, available_for_inference=False}>},
 'axis': Dim{'dyn-time:output'[?]},
 'n_out': <class 'returnn.util.basic.NotSpecified'>,
 'name': 'output',
 'network': <TFNetwork '' train=False>,
 'output': Data{'output_output', [T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}},
 'rec_step_by_step_opts': {'update_i_in_graph': False},
 'sources': [],
 'target': 'bpe',
 'unit': <SubnetworkRecCellSingleStep '/output(rec-subnet)'>}
Unhandled exception <class 'AssertionError'> in thread <_MainThread(MainThread, started 22947206223616)>, proc 10372.

Thread current, main, <_MainThread(MainThread, started 22947206223616)>:
(Excluded thread.)

That were all threads.
EXCEPTION
...
  File "/u/schmitt/src/returnn/returnn/tf/util/data.py", line 4318, in Data.placeholder
    line: self.sanity_check(assume_complete=False)
    locals:
      self = <local> Data{'concat_lm_att', [B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}
      self.sanity_check = <local> <bound method Data.sanity_check of Data{'concat_lm_att', [B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}>
      assume_complete = <not found>
  File "/u/schmitt/src/returnn/returnn/tf/util/data.py", line 2867, in Data.sanity_check
    line: assert self.placeholder.shape[i].value == self.batch_shape[i]
    locals:
      self = <local> Data{'concat_lm_att', [B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}
      self.placeholder = <local> <tf.Tensor 'output/rec/concat_lm_att/concat_sources/concat:0' shape=(?, ?) dtype=float32>
      self.placeholder.shape = <local> TensorShape([Dimension(None), Dimension(None)]), len = 2
      i = <local> 1
      value = <not found>
      self.batch_shape = <local> (None, 1024)
AssertionError
```

The full error log can also be seen in https://gist.github.com/robin-p-schmitt/3b3643efc7c8cfa8859f6d1a336549f9.

Using the same network dict during training and search does not cause any problems.





Discussion:
This shape looks wrong:
```
<MergeDimsLayer output/'att' out_type=Data{[B,F|F'((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_fw:feature)+((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_bw:feature)'[B]{ctx=loop('dyn-time:output'[?])}], ctx=loop('dyn-time:output'[?])}>]}
```
This comes from:
```
layer /output(rec-subnet)/'att0': [B,F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B],F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)] float32
layer /output(rec-subnet)/'att': [B,F|F'((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_fw:feature)+((energy:feature-dense)*lstm1_pool:conv:s0*lstm5_bw:feature)'[?]{ctx=loop('dyn-time:output'[?])}] float32
```
So, `merge_dims` on `except_batch` is wrong here. You want to merge only the static axes. So replace the `except_batch`. `except_batch` is almost never a good idea. Replace it e.g. by `["dim:1", "dim:2048"]`.

But there is also another problem. Why is there still the `T|'lstm1_pool:conv:s0'[B]` dim in `att0`?

This also looks wrong:
```
layer /output(rec-subnet)/'att_weights': ['lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0_rec_step_by_step'[B]] float32
```
Where does this `T|'lstm1_pool:conv:s0_rec_step_by_step'[B]` come from?

This already looks wrong here:
```
 - Data{'enc_ctx_output', [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],B,F|F'enc_ctx:feature-dense'(1024)]}
```
But the `enc_ctx` looks like:
```
layer /'enc_ctx': [T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)] float32
```
This is probably introduced via the base state var:
```
New state var 'base_value_enc_ctx': <tf.Variable 'IO/base_value_enc_ctx:0' shape=(?, ?, 1024) dtype=float32>, shape Data{'enc_ctx_output', [B,T|'lstm1_pool:conv:s0_rec_step_by_step'[B],F|F'enc_ctx:feature-dense'(1024)]}
```

I don't really have time to look into this myself now. Maybe you can take a look.
> I don't really have time to look into this myself now. Maybe you can take a look.

I also won't make it today but I will look into this on Monday.
> ```
> New state var 'base_value_enc_ctx': <tf.Variable 'IO/base_value_enc_ctx:0' shape=(?, ?, 1024) dtype=float32>, shape Data{'enc_ctx_output', [B,T|'lstm1_pool:conv:s0_rec_step_by_step'[B],F|F'enc_ctx:feature-dense'(1024)]}
> ```

Regarding the `'lstm1_pool:conv:s0_rec_step_by_step'[B]` dim: This is just the new name for the base state var of the dyn sizes of the base (enc ctx here, as another base state var). This represents the base dim. This is created here:
```
New state var 'base_size_lstm1_pool__conv__s0': <tf.Variable 'IO/base_size_lstm1_pool__conv__s0:0' shape=(?,) dtype=int32>, shape Data{'base_size_lstm1_pool__conv__s0', [B?], dtype='int32'}
```

Basically, everything inside the rec-step-by-step should always only see this new dim and not the original base dim (`'lstm1_pool:conv:s0'[B]`). Whenever we see `'lstm1_pool:conv:s0'[B]` later, this is the bug. It should be `'lstm1_pool:conv:s0_rec_step_by_step'[B]`.

So, this here looks wrong:
```
layer /output(rec-subnet)/'weight_feedback': [B,T|'lstm1_pool:conv:s0'[B],F|F'weight_feedback:feature-dense'(1024)] float32
```

> So, `merge_dims` on `except_batch` is wrong here. You want to merge only the static axes. So replace the `except_batch`. `except_batch` is almost never a good idea. Replace it e.g. by `["dim:1", "dim:2048"]`.

So, after fixing this, I now get the following error:

```
Create graph...
Loading network, train flag False, eval flag False, search flag False
layer /'data': [B,T|'time'[B],F|F'feature:data'(40)] float32
layer /'source_stddev': [B,T|'time'[B],F|F'feature:data'(40)] float32
layer /'source': [B,T|'time'[B],F|F'feature:data'(40)] float32
layer /'source0': [B,T|'time'[B],F'feature:data'(40),F|F'source0_split_dims1'(1)] float32
DEPRECATION WARNING: Explicitly specify in_spatial_dims when there is more than one spatial dim in the input.
This will be disallowed with behavior_version 8.
layer /'conv0': [B,T|'time'[B],F'feature:data'(40),F|F'conv0:channel'(32)] float32
layer /'conv0p': [B,T|'time'[B],'conv0p:conv:s1'(20),F|F'conv0:channel'(32)] float32
layer /'conv1': [B,T|'time'[B],'conv0p:conv:s1'(20),F|F'conv1:channel'(32)] float32
layer /'conv1p': [B,T|'time'[B],'conv1p:conv:s1'(10),F|F'conv1:channel'(32)] float32
layer /'conv_merged': [B,T|'time'[B],F|F'conv1p:conv:s1*conv1:channel'(320)] float32
DEPRECATION WARNING: MergeDimsLayer, only keep_order=True is allowed
This will be disallowed with behavior_version 6.
layer /'lstm0_fw': [T|'time'[B],B,F|F'lstm0_fw:feature'(1024)] float32
WARNING:tensorflow:From /u/schmitt/src/returnn/returnn/tf/util/basic.py:1408: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Call initializer instance with the dtype argument instead of passing it to the constructor
[2022-04-08 14:35:24,754] INFO: Run time: 0:00:20 CPU: 0.40% RSS: 930MB VMS: 12.95GB
layer /'lstm0_bw': [T|'time'[B],B,F|F'lstm0_bw:feature'(1024)] float32
layer /'lstm0_pool': [B,T|'lstm0_pool:conv:s0'[?],F|F'lstm0_fw:feature+lstm0_bw:feature'(2048)] float32
layer /'lstm1_fw': [T|'lstm0_pool:conv:s0'[B],B,F|F'lstm1_fw:feature'(1024)] float32
layer /'lstm1_bw': [T|'lstm0_pool:conv:s0'[B],B,F|F'lstm1_bw:feature'(1024)] float32
layer /'lstm1_pool': [B,T|'lstm1_pool:conv:s0'[?],F|F'lstm1_fw:feature+lstm1_bw:feature'(2048)] float32
layer /'lstm2_fw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm2_fw:feature'(1024)] float32
layer /'lstm2_bw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm2_bw:feature'(1024)] float32
layer /'lstm2_pool': [B,T|'lstm1_pool:conv:s0'[B],F|F'lstm2_fw:feature+lstm2_bw:feature'(2048)] float32
layer /'lstm3_fw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm3_fw:feature'(1024)] float32
layer /'lstm3_bw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm3_bw:feature'(1024)] float32
layer /'lstm3_pool': [B,T|'lstm1_pool:conv:s0'[B],F|F'lstm3_fw:feature+lstm3_bw:feature'(2048)] float32
layer /'lstm4_fw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm4_fw:feature'(1024)] float32
layer /'lstm4_bw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm4_bw:feature'(1024)] float32
layer /'lstm4_pool': [B,T|'lstm1_pool:conv:s0'[B],F|F'lstm4_fw:feature+lstm4_bw:feature'(2048)] float32
layer /'lstm5_fw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm5_fw:feature'(1024)] float32
layer /'lstm5_bw': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm5_bw:feature'(1024)] float32
layer /'encoder': [T|'lstm1_pool:conv:s0'[B],B,F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)] float32
layer /'ctc_out': [T|'lstm1_pool:conv:s0'[B],B,F|F'ctc_out:feature-dense'(1031)] float32
layer /'ctc_out_scores': [T|'lstm1_pool:conv:s0'[B],B,F|F'ctc_out:feature-dense'(1031)] float32
layer /'data:bpe': [B,T|'time:var:extern_data:bpe'[B]] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
layer /'ctc': [T|'lstm1_pool:conv:s0'[B],B,F|F'ctc_out:feature-dense'(1031)] float32
layer /'enc_value': [T|'lstm1_pool:conv:s0'[B],B,F'enc_value_split_dims0'(1),F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)] float32
layer /'enc_ctx': [T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)] float32
layer /'inv_fertility': [T|'lstm1_pool:conv:s0'[B],B,F|F'inv_fertility:feature-dense'(1)] float32
layer /'output': [T|'dyn-time:output'[?],B] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
New state var 'end_flag': <tf.Variable 'IO/end_flag:0' shape=(?,) dtype=bool>, shape Data{'end_flag', [B?], dtype='bool'}
New state var 'dyn_seq_len': <tf.Variable 'IO/dyn_seq_len:0' shape=(?,) dtype=int32>, shape Data{'dyn_seq_len', [B?], dtype='int32'}
New state var 'i': <tf.Variable 'IO/i:0' shape=() dtype=int32>, shape Data{'i', [], dtype='int32'}
Delayed: Layer 'target_embed' depends on choices {'output'}, deps on prev frame set()
New state var 'stochastic_var_scores_output': <tf.Variable 'IO/stochastic_var_scores_output:0' shape=(?, 1030) dtype=float32>, shape Data{'label_prob_output', [B,F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])}
New state var 'stochastic_var_choice_output': <tf.Variable 'IO/stochastic_var_choice_output:0' shape=(?,) dtype=int32>, shape Data{'output_output', [B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])}
New state var 'state_delayed/target_embed/output': <tf.Variable 'IO/state_delayed/target_embed/output:0' shape=(?, 621) dtype=float32>, shape Data{'target_embed_output', [B,F|F'target_embed:feature-dense'(621)], ctx=loop('dyn-time:output'[?])}
New state var 'state/att/output': <tf.Variable 'IO/state/att/output:0' shape=(?, 2048) dtype=float32>, shape Data{'att_output', [B,F|F'((energy:feature-dense)*lstm5_fw:feature)+((energy:feature-dense)*lstm5_bw:feature)'(2048)], ctx=loop('dyn-time:output'[?])}
New state var 'state/lm/extra/state/c': <tf.Variable 'IO/state/lm/extra/state/c:0' shape=(?, 1024) dtype=float32>, shape Data{'state/lm/extra/state/c', [B?,F|F'feature:state/lm/extra/state/c'(1024)]}
New state var 'state/lm/extra/state/h': <tf.Variable 'IO/state/lm/extra/state/h:0' shape=(?, 1024) dtype=float32>, shape Data{'state/lm/extra/state/h', [B?,F|F'feature:state/lm/extra/state/h'(1024)]}
New state var 'state/accum_att_weights/output': <tf.Variable 'IO/state/accum_att_weights/output:0' shape=(?, ?, 1) dtype=float32>, shape Data{'accum_att_weights_output', [B,T|'lstm1_pool:conv:s0'[B],F|F'inv_fertility:feature-dense'(1)], ctx=loop('dyn-time:output'[?])}
layer /output(rec-subnet)(delayed-update)/'data:bpe': [B] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
layer /output(rec-subnet)(delayed-update)/'output': [B] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
layer /output(rec-subnet)(delayed-update)/'target_embed': [B,F|F'target_embed:feature-dense'(621)] float32
New state var 'cond': <tf.Variable 'IO/cond:0' shape=() dtype=bool>, shape Data{'cond', [], dtype='bool'}
layer /output(rec-subnet)/'lm': [B,F|F'lm:feature'(1024)] float32
New state var 'base_size_lstm1_pool__conv__s0': <tf.Variable 'IO/base_size_lstm1_pool__conv__s0:0' shape=(?,) dtype=int32>, shape Data{'base_size_lstm1_pool__conv__s0', [B?], dtype='int32'}
New state var 'base_value_enc_value': <tf.Variable 'IO/base_value_enc_value:0' shape=(?, ?, 1, 2048) dtype=float32>, shape Data{'enc_value_output', [B,T|'lstm1_pool:conv:s0_rec_step_by_step'[B],F'enc_value_split_dims0'(1),F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)]}
New state var 'base_value_enc_ctx': <tf.Variable 'IO/base_value_enc_ctx:0' shape=(?, ?, 1024) dtype=float32>, shape Data{'enc_ctx_output', [B,T|'lstm1_pool:conv:s0_rec_step_by_step'[B],F|F'enc_ctx:feature-dense'(1024)]}
layer /output(rec-subnet)/'weight_feedback': [B,T|'lstm1_pool:conv:s0'[B],F|F'weight_feedback:feature-dense'(1024)] float32
layer /output(rec-subnet)/'att_query': [B,F|F'lm:feature'(1024)] float32
DEPRECATION WARNING: All inputs
 - Data{'enc_ctx_output', [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],B,F|F'enc_ctx:feature-dense'(1024)]}
 - Data{'weight_feedback_output', [B,T|'lstm1_pool:conv:s0'[B],F|F'weight_feedback:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])}
 - Data{'att_query_output', [B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])}
require broadcasting to 
  Data{'att_energy_in_output', [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)]}.
This must be explicitly allowed, e.g. by specifying out_shape.
This will be disallowed with behavior_version 4.
layer /output(rec-subnet)/'att_energy_in': [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)] float32
layer /output(rec-subnet)/'energy_tanh': [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)] float32
layer /output(rec-subnet)/'energy': [T|'lstm1_pool:conv:s0_rec_step_by_step'[B],'lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1)] float32
layer /output(rec-subnet)/'att_weights0': ['lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0_rec_step_by_step'[B]] float32
layer /output(rec-subnet)/'att_weights': ['lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0_rec_step_by_step'[B]] float32
layer /output(rec-subnet)/'att0': [B,F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B],F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)] float32
layer /output(rec-subnet)/'att': [B,T|'lstm1_pool:conv:s0'[B],F|F'((energy:feature-dense)*lstm5_fw:feature)+((energy:feature-dense)*lstm5_bw:feature)'(2048)] float32
layer /output(rec-subnet)/'readout_in': [B,T|'lstm1_pool:conv:s0'[B],F|F'readout_in:feature-dense'(1000)] float32
layer /output(rec-subnet)/'readout': [B,T|'lstm1_pool:conv:s0'[B],F|F'(readout_in:feature-dense)//2'(500)] float32
layer /output(rec-subnet)/'data:bpe': [B] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
layer /output(rec-subnet)/'label_prob': [B,T|'lstm1_pool:conv:s0'[B],F|F'bpe:sparse-dim'(1030)] float32
layer /output(rec-subnet)/'output': [B] int32 sparse_dim=Dim{F'bpe:sparse-dim'(1030)}
Exception creating layer /output(rec-subnet)/'output' of class ChoiceStateVarLayer with opts:
{'_name': 'output',
 '_network': <TFNetwork '/output(rec-subnet)' parent_layer=<RecStepByStepLayer 'output' out_type=Data{[T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}}> train=False>,
 '_target_layers': {'bpe': <DataNotAvailableLayer output/'data:bpe' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, available_for_inference=False, ctx=loop('dyn-time:output'[?])}>},
 'beam_size': 12,
 'initial_output': 0,
 'n_out': 1030,
 'name': 'output',
 'network': <TFNetwork '/output(rec-subnet)' parent_layer=<RecStepByStepLayer 'output' out_type=Data{[T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}}> train=False>,
 'output': Data{'output_output', [B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])},
 'sources': [<SoftmaxLayer output/'label_prob' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])}>],
 'target': ['bpe']}

ERROR: Got exception during in-loop construction of layer 'output':
AssertionError: 

Template network (check out types / shapes):
output: <_TemplateLayer(ChoiceStateVarLayer)(:template:choice_state_var) output/'output' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])} (construction stack None)>
label_prob: <_TemplateLayer(SoftmaxLayer)(:template:softmax) output/'label_prob' out_type=Data{[B,F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])} (construction stack 'output')>
readout: <_TemplateLayer(ReduceOutLayer)(:template:reduce_out) output/'readout' out_type=Data{[B,F|F'(readout_in:feature-dense)//2'(500)], ctx=loop('dyn-time:output'[?])} (construction stack 'label_prob')>
readout_in: <_TemplateLayer(LinearLayer)(:template:linear) output/'readout_in' out_type=Data{[B,F|F'readout_in:feature-dense'(1000)], ctx=loop('dyn-time:output'[?])} (construction stack 'readout')>
lm: <_TemplateLayer(RecLayer)(:template:rec) output/'lm' out_type=Data{[B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'readout_in')>
target_embed: <_TemplateLayer(LinearLayer)(:template:linear) output/'target_embed' out_type=Data{[B,F|F'target_embed:feature-dense'(621)], ctx=loop('dyn-time:output'[?])} (construction stack 'lm')>
att: <_TemplateLayer(MergeDimsLayer)(:template:merge_dims) output/'att' out_type=Data{[B,F|F'((energy:feature-dense)*lstm5_fw:feature)+((energy:feature-dense)*lstm5_bw:feature)'(2048)], ctx=loop('dyn-time:output'[?])} (construction stack 'lm')>
att0: <_TemplateLayer(GenericAttentionLayer)(:template:generic_attention) output/'att0' out_type=Data{[B,F'energy:feature-dense'(1),F|F'lstm5_fw:feature+lstm5_bw:feature'(2048)], ctx=loop('dyn-time:output'[?])} (construction stack 'att')>
att_weights: <_TemplateLayer(DropoutLayer)(:template:dropout) output/'att_weights' out_type=Data{[B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B]], ctx=loop('dyn-time:output'[?])} (construction stack 'att0')>
att_weights0: <_TemplateLayer(SoftmaxOverSpatialLayer)(:template:softmax_over_spatial) output/'att_weights0' out_type=Data{[B,F|F'energy:feature-dense'(1),T|'lstm1_pool:conv:s0'[B]], ctx=loop('dyn-time:output'[?])} (construction stack 'att_weights')>
energy: <_TemplateLayer(LinearLayer)(:template:linear) output/'energy' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'energy:feature-dense'(1)], ctx=loop('dyn-time:output'[?])} (construction stack 'att_weights0')>
energy_tanh: <_TemplateLayer(ActivationLayer)(:template:activation) output/'energy_tanh' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy')>
att_energy_in: <_TemplateLayer(CombineLayer)(:template:combine) output/'att_energy_in' out_type=Data{[T|'lstm1_pool:conv:s0'[B],B,F|F'enc_ctx:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'energy_tanh')>
weight_feedback: <_TemplateLayer(LinearLayer)(:template:linear) output/'weight_feedback' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'weight_feedback:feature-dense'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'att_energy_in')>
accum_att_weights: <_TemplateLayer(EvalLayer)(:template:eval) output/'accum_att_weights' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'inv_fertility:feature-dense'(1)], ctx=loop('dyn-time:output'[?])} (construction stack 'weight_feedback')>
att_query: <_TemplateLayer(LinearLayer)(:template:linear) output/'att_query' out_type=Data{[B,F|F'lm:feature'(1024)], ctx=loop('dyn-time:output'[?])} (construction stack 'att_energy_in')>
data:bpe: <_TemplateLayer(SourceLayer)(:template:source) output/'data:bpe' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, available_for_inference=False, ctx=loop('dyn-time:output'[?])} (construction stack 'label_prob')>
end: <_TemplateLayer(CompareLayer)(:template:compare) output/'end' out_type=Data{[B], dtype='bool', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])} (construction stack None)>

Exception creating layer /'output' of class RecStepByStepLayer with opts:
{'_name': 'output',
 '_network': <TFNetwork '' train=False>,
 '_orig_sources': [],
 '_target_layers': {'bpe': <DataNotAvailableLayer 'data:bpe' out_type=Data{[B,T|'time:var:extern_data:bpe'[B]], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, available_for_inference=False}>},
 'axis': Dim{'dyn-time:output'[?]},
 'n_out': <class 'returnn.util.basic.NotSpecified'>,
 'name': 'output',
 'network': <TFNetwork '' train=False>,
 'output': Data{'output_output', [T|'dyn-time:output'[?],B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}},
 'rec_step_by_step_opts': {'update_i_in_graph': False},
 'sources': [],
 'target': 'bpe',
 'unit': <SubnetworkRecCellSingleStep '/output(rec-subnet)'>}
Unhandled exception <class 'AssertionError'> in thread <_MainThread(MainThread, started 23260537734912)>, proc 28980.

Thread current, main, <_MainThread(MainThread, started 23260537734912)>:
(Excluded thread.)

That were all threads.
EXCEPTION
...
  File "/u/schmitt/src/returnn/tools/compile_tf_graph.py", line 1407, in ChoiceStateVarLayer.__init__
    line: self.stochastic_var.assign_score(source)
    locals:
      self = <local> <ChoiceStateVarLayer output/'output' out_type=Data{[B], dtype='int32', sparse_dim=Dim{F'bpe:sparse-dim'(1030)}, ctx=loop('dyn-time:output'[?])}>
      self.stochastic_var = <local> <__main__.RecStepByStepLayer.StochasticVar object at 0x15273ca22f70>
      self.stochastic_var.assign_score = <local> <bound method RecStepByStepLayer.StochasticVar.assign_score of <__main__.RecStepByStepLayer.StochasticVar object at 0x15273ca22f70>>
      source = <local> <SoftmaxLayer output/'label_prob' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])}>
  File "/u/schmitt/src/returnn/tools/compile_tf_graph.py", line 1329, in RecStepByStepLayer.StochasticVar.assign_score
    line: assert source.output.is_batch_major and len(source.output.shape) == 1
    locals:
      source = <local> <SoftmaxLayer output/'label_prob' out_type=Data{[B,T|'lstm1_pool:conv:s0'[B],F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])}>
      source.output = <local> Data{'label_prob_output', [B,T|'lstm1_pool:conv:s0'[B],F|F'bpe:sparse-dim'(1030)], ctx=loop('dyn-time:output'[?])}
      source.output.is_batch_major = <local> True
      len = <builtin> <built-in function len>
      source.output.shape = <local> (None, 1030)
AssertionError
```

And I think this is the error which you are talking about, with `'lstm1_pool:conv:s0'[B]`.

The corresponding layer is this here:

```
"label_prob": {
                "class": "softmax",
                "dropout": 0.3,
                "from": ["readout"],
                "loss": "ce",
                "loss_opts": {"label_smoothing": 0.1},
                "target": "bpe",
            },
```

The full network and log can be found here: https://gist.github.com/robin-p-schmitt/76808f4a70d12a73f6a1616c1f3a64ab
Yes I think I know the problem. I'm currently fixing it.

Resolving patch:
diff --git a/tools/compile_tf_graph.py b/tools/compile_tf_graph.py
--- a/tools/compile_tf_graph.py
+++ b/tools/compile_tf_graph.py
@@ -618,6 +618,25 @@ def _construct(self, prev_outputs, prev_extra, i, data=None,
       data_ = self.net.extern_data.data[key]
       data[key] = self._tiled(data_, value)
 
+    # Fixup all base state var dim tags in the template layers.
+    # These templates are used to define prev layers,
+    # e.g. like prev accum att weights, so we might need the parent dyn sizes now.
+    all_dyn_parent_dim_tags = set()
+    for layer in self.parent_net.get_root_network().get_all_layers_deep():
+      if self.net in layer.network.get_network_hierarchy():
+        continue
+      all_dyn_parent_dim_tags.update([
+        dim for dim in layer.output.dim_tags
+        if dim.dimension is None and not dim.is_batch_dim()])
+    for layer in self.layer_data_templates.values():
+      assert isinstance(layer, _TemplateLayer)
+      dim_tags = list(layer.output.dim_tags)
+      for _i, tag in enumerate(dim_tags):
+        if tag in all_dyn_parent_dim_tags:
+          dim_tags[_i] = self.get_parent_dim_tag(tag)
+      layer.output = layer.output.copy_template_new_dim_tags(dim_tags)
+      layer.kwargs["output"] = layer.output
+
     super(SubnetworkRecCellSingleStep, self)._construct(
       prev_outputs=prev_outputs, prev_extra=prev_extra, i=i, data=data,
       inputs_moved_out_tas=inputs_moved_out_tas, needed_outputs=needed_outputs)


Test patch:
diff --git a/tests/test_tools.py b/tests/test_tools.py
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -46,7 +46,7 @@ def run(*args):
 # Tests for compile_tf_graph.py
 ###############################
 
-rec_encoder_decoder_config = """
+rec_encoder_decoder_simple_config = """
 #!rnn.py
 network = {
     "enc0": {"class": "linear", "from": "data", "activation": "sigmoid", "n_out": 3},
@@ -69,6 +69,42 @@ def run(*args):
 """
 
 
+rec_encoder_decoder_att_config = """
+#!rnn.py
+network = {
+    "encoder": {"class": "linear", "from": "data", "activation": "sigmoid", "n_out": 3},
+    "enc_ctx": {"class": "linear", "from": "encoder", "n_out": 10},
+    "inv_fertility": {"class": "linear", "activation": "sigmoid", "with_bias": False, "from": "encoder", "n_out": 1},
+    "output": {
+      "class": "rec", "from": [], "target": "classes",
+      "unit": {
+        "weight_feedback": {"class": "linear", "activation": None, "with_bias": False, "from": "prev:accum_att_weights", "n_out": 10},
+        "prev_s_transformed": {"class": "linear", "activation": None, "with_bias": False, "from": "prev:s", "n_out": 10},
+        "energy_in": {"class": "combine", "kind": "add", "from": ["base:enc_ctx", "weight_feedback", "prev_s_transformed"], "n_out": 10},
+        "energy_tanh": {"class": "activation", "activation": "tanh", "from": "energy_in"},
+        "energy": {"class": "linear", "activation": None, "with_bias": False, "from": "energy_tanh", "n_out": 1},
+        "att_weights": {"class": "softmax_over_spatial", "from": "energy"},
+        "accum_att_weights": {
+          "class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
+          "eval": "source(0) + source(1) * source(2) * 0.5",
+          "out_type": {"dim": 1, "shape": (None, 1)}},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
+
+        "s": {"class": "rec", "unit": "lstm", "from": ["att", "prev:embed"], "n_out": 3},
+        "prob": {"class": "softmax", "from": "s", "loss": "ce", "target": "classes"},
+        "output": {"class": "choice", "beam_size": 4, "from": "prob", "target": "classes", "initial_output": 0},
+        "embed": {"class": "linear", "from": "prev:output", "activation": "sigmoid", "n_out": 3},
+        "end": {"class": "compare", "from": "output", "value": 0}
+      }
+    },
+    "decision": {"class": "decide", "from": "output", "loss": "edit_distance"}
+}
+num_inputs = 5
+num_outputs = 3
+use_tensorflow = True
+"""
+
+
 rec_transducer_time_sync_config = """
 #!rnn.py
 network = {
@@ -113,7 +149,7 @@ def run(*args):
 def test_compile_tf_graph_basic():
   tmp_dir = tempfile.mkdtemp()
   with open(os.path.join(tmp_dir, "returnn.config"), "wt") as config:
-    config.write(rec_encoder_decoder_config)
+    config.write(rec_encoder_decoder_simple_config)
   args = [
     "tools/compile_tf_graph.py",
     "--output_file",
@@ -130,10 +166,26 @@ def test_compile_tf_graph_basic_second_run():
   test_compile_tf_graph_basic()
 
 
-def test_compile_tf_graph_enc_dec_recurrent_step():
+def test_compile_tf_graph_enc_dec_simple_recurrent_step():
+  tmp_dir = tempfile.mkdtemp()
+  with open(os.path.join(tmp_dir, "returnn.config"), "wt") as config:
+    config.write(rec_encoder_decoder_simple_config)
+  args = [
+    "tools/compile_tf_graph.py",
+    "--output_file",
+    os.path.join(tmp_dir, "graph.metatxt"),
+    "--rec_step_by_step",
+    "output",
+    os.path.join(tmp_dir, "returnn.config")
+  ]
+  run(*args)
+
+
+def test_compile_tf_graph_enc_dec_att_recurrent_step():
+  # https://github.com/rwth-i6/returnn/issues/1016
   tmp_dir = tempfile.mkdtemp()
   with open(os.path.join(tmp_dir, "returnn.config"), "wt") as config:
-    config.write(rec_encoder_decoder_config)
+    config.write(rec_encoder_decoder_att_config)
   args = [
     "tools/compile_tf_graph.py",
     "--output_file",