Getting batch dim tags fails when using _rel_shift function for relative pos embeddings

[mmz33]

I am getting this issue when running python3 -m tests.test_nn_conformer

This is the error backtrace:

EXCEPTION
Traceback (most recent call last):
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.7/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    line: return _run_code(code, main_globals, None,
                           "__main__", mod_spec)
    locals:
      _run_code = <global> <function _run_code at 0x10a6caa60>
      code = <local> <code object <module> at 0x10a689930, file "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py", line 3>
      main_globals = <local> {'__name__': '__main__', '__doc__': '\nTest nn.conformer.\n', '__package__': 'returnn_common.tests', '__loader__': <_frozen_importlib_external.SourceFileLoader object at 0x10a737240>, '__spec__': ModuleSpec(name='tests.test_nn_conformer', loader=<_frozen_importlib_external.SourceFileLoader object..., len = 10
      mod_spec = <local> ModuleSpec(name='tests.test_nn_conformer', loader=<_frozen_importlib_external.SourceFileLoader object at 0x10a737240>, origin='/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py')
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.7/lib/python3.7/runpy.py", line 85, in _run_code
    line: exec(code, run_globals)
    locals:
      exec = <builtin> <built-in function exec>
      code = <local> <code object <module> at 0x10a689930, file "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py", line 3>
      run_globals = <local> {'__name__': '__main__', '__doc__': '\nTest nn.conformer.\n', '__package__': 'returnn_common.tests', '__loader__': <_frozen_importlib_external.SourceFileLoader object at 0x10a737240>, '__spec__': ModuleSpec(name='tests.test_nn_conformer', loader=<_frozen_importlib_external.SourceFileLoader object..., len = 10
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py", line 7, in <module>
    line: from . import _setup_test_env  # noqa
    locals:
      from = <not found>
      from.import = <not found>
      _setup_test_env = <not found>
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/_setup_test_env.py", line 307, in <module>
    line: setup()
    locals:
      setup = <local> <function setup at 0x10a71f158>
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/_setup_test_env.py", line 134, in setup
    line: _main(main_mod_)
    locals:
      _main = <global> <function _main at 0x10a71f6a8>
      main_mod_ = <local> <module 'returnn_common.tests.test_nn_conformer' from '/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py'>
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/_setup_test_env.py", line 298, in _main
    line: value()
    locals:
      value = <local> <function test_nn_conformer at 0x1585daa60>
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/tests/test_nn_conformer.py", line 39, in test_nn_conformer
    line: out, _ = conformer(data, in_spatial_dim=time_dim)
    locals:
      out = <not found>
      _ = <not found>
      conformer = <local> <ConformerEncoder>
      data = <local> <Tensor /'data:data' [B,T|'time'[B],F|F'input'(10)]>
      in_spatial_dim = <not found>
      time_dim = <local> Dim{'time'[B]}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/conformer.py", line 278, in ConformerEncoder.__call__
    line: x = self.layers(x, axis=out_spatial_dim)
    locals:
      x = <local> <Tensor /'conformer_encoder'/'dropout' [B,T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|F'out'(14)] via 'dropout'>
      self = <local> <ConformerEncoder>
      self.layers = <local> <Sequential>, len = 2
      axis = <not found>
      out_spatial_dim = <local> Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B]}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/container.py", line 80, in Sequential.__call__
    line: inp = module(inp, **kwargs)
    locals:
      inp = <local> <Tensor /'conformer_encoder'/'dropout' [B,T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|F'out'(14)] via 'dropout'>
      module = <local> <ConformerEncoderLayer>
      kwargs = <local> {'axis': Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B]}}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/conformer.py", line 203, in ConformerEncoderLayer.__call__
    line: x_mhsa = self.self_att(x_mhsa_ln, axis=axis)
    locals:
      x_mhsa = <not found>
      self = <local> <ConformerEncoderLayer>
      self.self_att = <local> <RelPosSelfAttention>
      x_mhsa_ln = <local> <Tensor /'conformer_encoder'/'layers'/'0'/'self_att_layer_norm'/'add_0'/'combine'/'add' [B,T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|F'out'(14)] via 'combine'>
      axis = <local> Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B]}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/attention.py", line 204, in RelPosSelfAttention.__call__
    line: matrix_bd = self._rel_shift(matrix_bd, axis, pos_emb_spatial_dim, hist_dim)
    locals:
      matrix_bd = <local> <Tensor /'conformer_encoder'/'layers'/'0'/'self_att'/'dot_0' [T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|'num_heads'(2),B] via 'dot'>
      self = <local> <RelPosSelfAttention>
      self._rel_shift = <local> <bound method RelPosSelfAttention._rel_shift of <class 'returnn_common.nn.attention.RelPosSelfAttention'>>
      axis = <local> Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B]}
      pos_emb_spatial_dim = <local> Dim{'(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)+-1+(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)'[B]}
      hist_dim = <local> Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0:kv'[?]}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/attention.py", line 226, in RelPosSelfAttention._rel_shift
    line: batch_dims = x.batch_dims_ordered((axis, pos_emb_spatial_dim))
    locals:
      batch_dims = <not found>
      x = <local> <Tensor /'conformer_encoder'/'layers'/'0'/'self_att'/'dot_0' [T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|'num_heads'(2),B] via 'dot'>
      x.batch_dims_ordered = <local> <bound method Tensor.batch_dims_ordered of <Tensor /'conformer_encoder'/'layers'/'0'/'self_att'/'dot_0' [T|'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B],F|'num_heads'(2),B] via 'dot'>>
      axis = <local> Dim{'conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0'[B]}
      pos_emb_spatial_dim = <local> Dim{'(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)+-1+(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)'[B]}
  File "/Users/mzeineldeen/Desktop/dev/returnn_common/nn/base.py", line 201, in Tensor.batch_dims_ordered
    line: batch_dims.remove(remove_)
    locals:
      batch_dims = <local> [Dim{'num_heads'(2)}, Dim{B}]
      batch_dims.remove = <local> <built-in method remove of list object at 0x1597f1308>
      remove_ = <local> Dim{'(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)+-1+(conformer_encoder/conv_subsample_layer/pool2d_0/_pool_nd:out-spatial-dim0)'[B]}
ValueError: list.remove(x): x not in list

For some reason the input x does not have correct dim tags right? this should be matrix_bd which has shape (batch, head, time1, 2*time1-1) so the time dimension tags are missing.

[albertz]

Yup. This is exactly this RETURNN issue: rwth-i6/returnn#1154

[albertz]

Ok, should be fixed with new RETURNN version.