Repository: rwth-i6/returnn

Issue:
2D/3D convolution/pooling (`ConvLayer`/`PoolLayer`/`TransposedConvLayer`) depends on order of axes
We want to avoid anything which depends on the order of axes.



Discussion:
Some related discussion is in #597. E.g. there was the proposal of `in_to_out_dims`. E.g. for `ConvLayer` with 2D conv, it would have 3 entries: 2 for the two spatial dims, and then one for the linear feature transformation. For `PoolLayer` and others which do not do a linear transformation, it would contain only entries for the spatial dims.

The `filter_size` and related options would not be a list/tuple but instead a `dict[DimensionTag, int]` which also specifies the spatial axes it operates on.

One open question is the native format (shape, order of axes) of the `tf.Variable`. For that, the order of spatial axes matters.

Also, the native order of the input could change for various reasons, so we must not make it dependent on that.

This would be a very welcome addition, I remember I had some cases where I had to play around a lot to make Returnn run the filters over the axes I wanted... Isn't it fine to have a logic that moves the axes in the optimal order? Just like the [T,B,F] transformation for the nativelstm? I am not sure what is happening right now, but I think a re-ordering based in the selected Dimension tags is possible, and the tf.Variable can be defined accordingly.
One [core principle of RETURNN](https://github.com/rwth-i6/returnn/wiki/RETURNN-principles) is that the order of the axes should never ever matter in any way. This core principles allows layers to reorder axes for efficiency (e.g. like `RecLayer`). And this is exactly what this issue here is about, because this is violated by `ConvLayer` for 2D/3D convolution, and similarly for `PoolLayer`, and some other similar layers.

The proposal here would also not have any order. It would just be like `filter_size={dimtag1: 3, dimtag2: 5}` or so.

But then, the order of axes for the underlying `tf.Variable` is ambiguous.

> Isn't it fine to have a logic that moves the axes in the optimal order

We already do that. But this should not matter to the user.

> I had to play around a lot to make Returnn run the filters over the axes I wanted

This is not so much this issue here. This issue here is about the order of spatial axes, which matters, although it should not.

I assume in your case, you had some mixup between spatial and feature dim. You do not want to reorder the axes. You simply need to redefine which axis is the feature dim (channel dim). You can use `ReinterpretDataLayer` with `set_axes: {F: ...}` for that.

> But then, the order of axes for the underlying `tf.Variable` is ambiguous.

Ah, you mean for the cases where there is no "more efficient" order ? like it is not important if h or w comes first for images? Well... then just a heuristic is fine or not? Like sort by largest filter size first, or sort by name of dimension tag or so...
> > But then, the order of axes for the underlying `tf.Variable` is ambiguous.
> 
> Ah, you mean for the cases where there is no "more efficient" order ? like it is not important if h or w comes first for images? Well... then just a heuristic is fine or not? Like sort by largest filter size first, or sort by name of dimension tag or so...

Yea, sth like that. I'm sure there are some possible solutions. Although I'm not so happy with all the potential solutions I thought about so far.

I don't really like heuristics. And what when the filter sizes are the same (which is not so uncommon)?

I also don't really know which order of the filter dims is optimal. Maybe sorting is actually good, and largest first makes it most efficient. Who knows. Probably depends also on the hardware. Or even CUDA or cuDNN version. E.g. it was true for long time that batch-feature-major was more efficient on GPU than batch-spatial-major, but I recently read that some of the more recent professional Nvidia GPUs probably make batch-spatial-major more efficient.

The current best way I can think of is to use heuristics anyway, or maybe just use the input spatial axis order, but then also store the order explicitly somewhere in the TF checkpoint. And when we load the checkpoint, make sure the order is fine, or maybe reorder accordingly.

The dim tag names are also not really good for identification. Esp the internal names (`description` actually...) are not really fixed and are being changed from time to time. See also #634. But when we make new dim tags more explicit by the user (#597), i.e. no internal names are used, then this should be less ambiguous.

However, when you import some model, and the dim tag names do not match for whatever reason, should the import fail then?

> However, when you import some model, and the dim tag names do not match for whatever reason, should the import fail then?

Oh, I see... yes then this is really not easy to choose...
One option is to make the order explicit in the config. So `filter_size=OrderedDict(...)`.
Or, another option (from the comment [here](https://github.com/rwth-i6/returnn/issues/597#issuecomment-927183267)) is to have an argument `in_spatial_dims` (list/tuple of dim tags / axes) which defines the order. Then `filter_size` and all the other arguments can stay as they are.
`in_spatial_dims` is implemented now for `ConvLayer` and `PoolLayer` (#789).

However, it is not mandatory yet. Maybe this should be done as well, via new behavior version.
`TransposedConvLayer` also has the same problems. Fixed via #791.


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
@@ -4089,8 +4089,8 @@ def __init__(self, switch_axes=None, size_base=None, set_axes=None,
     :param dict[str,int|str] set_axes:
       This can be used to overwrite the special axes like time_dim_axis or feature_dim_axis.
       For that, use keys "B","T" or "F", and a value via :func:`Data.get_axis_from_description`.
-    :param dict[str|DimensionTag,DimensionTag] set_dim_tags: axis -> new dim tag. assigns new dim tags.
-      If the dim tag is yet undefined, this will not use same_dim_tags_as (declare_same_as)
+    :param dict[str|DimensionTag,DimensionTag]|None set_dim_tags: axis -> new dim tag. assigns new dim tags.
+      If the passed dim tag is yet undefined, this will not use same_dim_tags_as (declare_same_as)
       but create a new dim tag.
       This option is useful for generalized self attention (https://github.com/rwth-i6/returnn/issues/391).
     :param bool enforce_batch_major:
@@ -4163,7 +4163,7 @@ def get_out_data_from_opts(cls, name, sources,
     :param str|list[str] switch_axes: e.g. "bt" to switch batch and time axes
     :param LayerBase|None size_base: similar as size_target
     :param dict[str,int] set_axes:
-    :param dict[str|DimensionTag,DimensionTag] set_dim_tags:
+    :param dict[str|DimensionTag,DimensionTag]|None set_dim_tags:
     :param bool enforce_batch_major:
     :param bool enforce_time_major:
     :param bool|None set_sparse: if bool, set sparse value to this
@@ -4276,7 +4276,7 @@ def __init__(self, filter_size, padding, strides=1, dilation_rate=1, groups=1,
     :param int|tuple[int] dilation_rate: dilation for the spatial dims
     :param int groups: grouped convolution
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :param int|None n_out: number of outgoing features
     :param DimensionTag|None out_dim:
     :param list[DimensionTag]|None out_spatial_dims:
@@ -4327,7 +4327,7 @@ def __init__(self, filter_size, padding, strides=1, dilation_rate=1, groups=1,
       input_data = input_data.copy_with_feature_dim_axis(-1)
       in_spatial_dims_ = input_data.dim_tags[num_batch_dims:-1]
     if in_spatial_dims:
-      assert list(in_spatial_dims_) == list(in_spatial_dims)
+      assert list(in_spatial_dims_) == [input_data.get_dim_tag_from_description(d) for d in in_spatial_dims]
     assert len(in_spatial_dims_) == len(filter_size), (
       "%s: in_spatial_dims %s not matching to filter_size %s" % (self, in_spatial_dims_, filter_size))
     assert input_data.batch_ndim - num_batch_dims - 1 == len(filter_size), (
@@ -4472,6 +4472,17 @@ def set_output_dim_tags(cls, output, num_batch_dims, in_spatial_dims, out_spatia
         else:
           out_tag.declare_same_as(size_tag)
 
+  @classmethod
+  def _check_defined_in_spatial_dims(cls, cond):
+    """
+    :param bool cond:
+    """
+    from returnn.util import BehaviorVersion
+    BehaviorVersion.require(
+      condition=cond,
+      message="Explicitly specify in_spatial_dims when there is more than one spatial dim in the input.",
+      version=8)
+
   @classmethod
   def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
                       input_expand_dims=0, input_split_feature_dim=None, input_add_feature_dim=False):
@@ -4479,7 +4490,7 @@ def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
     :param Data input_data:
     :param returnn.tf.network.TFNetwork network:
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :param int input_expand_dims: number of spatial dims to add to the input
     :param None|int input_split_feature_dim: if set, like input_add_feature_dim it will add a new feature dim
       which is of value input_split_feature_dim, and the original input feature dim
@@ -4491,14 +4502,26 @@ def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
     """
     assert not input_data.sparse
     assert input_data.have_batch_axis()
+    if input_expand_dims or input_split_feature_dim or input_add_feature_dim:
+      if not in_spatial_dims:
+        # Set it here to allow to handle the extension logic below.
+        if input_data.placeholder is None and not input_data.have_feature_axis():
+          # This is still a template. Ignore errors and make it valid.
+          input_data = input_data.copy_add_feature_dim()
+        in_spatial_dims = [input_data.dim_tags[a] for a in input_data.get_spatial_batch_axes()]
+        cls._check_defined_in_spatial_dims(len(in_spatial_dims) == 1)
     if input_expand_dims:
       for i in range(input_expand_dims):
-        input_data = input_data.copy_add_spatial_dim()
+        dim_tag = DimensionTag(kind=DimensionTag.Types.Spatial, description="input_expand_dims:%i" % i, dimension=1)
+        input_data = input_data.copy_add_dim_by_tag(dim_tag, unbroadcast=True)
+        in_spatial_dims.append(dim_tag)
     if input_split_feature_dim:
       # Split the feature dimension.
       input_data = input_data.copy_split_feature_dim(input_split_feature_dim)
+      in_spatial_dims.append(input_data.dim_tags[input_data.feature_dim_axis - 1])
     if input_add_feature_dim:
-      # Add a feature dimension; any other static dims will be used as dynamic dims below.
+      # Add a new feature dimension; the old feature dim becomes a spatial dim.
+      in_spatial_dims.append(input_data.dim_tags[input_data.feature_dim_axis])
       input_data = input_data.copy_add_feature_dim()
     if in_dim:
       assert isinstance(in_dim, DimensionTag)
@@ -4506,8 +4529,9 @@ def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
         input_data.feature_dim_axis = input_data.get_axis_from_description(in_dim)
     if in_spatial_dims:
       from returnn.tf.util.data import BatchDim
-      assert all(isinstance(d, DimensionTag) for d in in_spatial_dims)
+      assert all(isinstance(d, (DimensionTag, str)) for d in in_spatial_dims)
       axes = [input_data.get_axis_from_description(d) for d in in_spatial_dims]  # also requires unique dim tags
+      in_spatial_dims = [input_data.dim_tags[a] for a in axes]
       assert sorted(set(axes)) == sorted(axes), (
         "in_spatial_dims %s must be unique but map to %s" % (in_spatial_dims, axes))
       if sorted(axes) != axes:
@@ -4518,7 +4542,7 @@ def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
           input_data = input_data.copy_move_axis(old_axis=a, new_axis=first + i)
         axes = [input_data.get_axis_from_description(d) for d in in_spatial_dims]
         assert sorted(axes) == axes
-      assert input_data.feature_dim_axis not in axes
+      assert input_data.feature_dim_axis not in axes, "invalid in_spatial_dims %s" % (in_spatial_dims,)
       expected_dims = {BatchDim, input_data.feature_dim_or_sparse_dim} | set(in_spatial_dims)
       assert len(expected_dims) == 2 + len(in_spatial_dims)
       # There might be more dims in the input than we expect.
@@ -4548,6 +4572,10 @@ def transform_input(cls, input_data, network, in_dim=None, in_spatial_dims=None,
           input_data = input_data.copy_move_axis(old_axis=a, new_axis=batch_axis_idx)
         batch_axis_idx += 1
       num_batch_dims = batch_axis_idx
+      in_spatial_dims = [
+        d for (i, d) in enumerate(input_data.dim_tags)
+        if i >= num_batch_dims and i != input_data.feature_dim_axis]
+      cls._check_defined_in_spatial_dims(len(in_spatial_dims) <= 1)
     return input_data, num_batch_dims
 
   @classmethod
@@ -4604,7 +4632,7 @@ def get_out_data_from_opts(
     :param bool input_add_feature_dim:
     :param None|int input_split_feature_dim:
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :param int|None n_out: number of outgoing features
     :param DimensionTag|None out_dim:
     :param list[DimensionTag]|None out_spatial_dims:
@@ -4719,7 +4747,7 @@ def __init__(self, mode, pool_size, padding="VALID", dilation_rate=1, strides=No
     :param tuple[int]|int dilation_rate:
     :param tuple[int]|int|None strides: in contrast to tf.nn.pool, the default (if it is None) will be set to pool_size
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :param DimensionTag|None out_dim:
     :param list[DimensionTag]|None out_spatial_dims:
     :param bool use_channel_first: if set, will transform input to NCHW format
@@ -4766,7 +4794,7 @@ def __init__(self, mode, pool_size, padding="VALID", dilation_rate=1, strides=No
       input_data = input_data.copy_with_feature_dim_axis(-1)
       in_spatial_dims_ = input_data.dim_tags[num_batch_dims:-1]
     if in_spatial_dims:
-      assert list(in_spatial_dims_) == list(in_spatial_dims)
+      assert list(in_spatial_dims_) == [input_data.get_dim_tag_from_description(d) for d in in_spatial_dims]
     assert len(in_spatial_dims_) == len(pool_size)
     assert input_data.batch_ndim - num_batch_dims - 1 == len(pool_size), (
       "%s: pool-size-dimension does not match the input data. " % self +
@@ -4813,7 +4841,7 @@ def get_out_data_from_opts(cls, name, sources, network,
     :param int|tuple[int]|list[int] dilation_rate:
     :param str padding:
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :param DimensionTag|None out_dim:
     :param list[DimensionTag]|None out_spatial_dims:
     :param bool use_channel_first:
@@ -4915,6 +4943,10 @@ def __init__(self, filter_size, strides=None,
     :param str padding: "same" or "valid"
     :param list[int]|int remove_padding:
     :param list[int|None]|int|None output_padding:
+    :param DimensionTag|None in_dim:
+    :param list[DimensionTag|str]|None in_spatial_dims:
+    :param DimensionTag|None out_dim:
+    :param list[DimensionTag]|None out_spatial_dims:
     :param bool with_bias: whether to add a bias. enabled by default.
       Note that the default is different from ConvLayer!
     :param str|None activation:
@@ -4937,6 +4969,10 @@ def __init__(self, filter_size, strides=None,
       in_dim=in_dim, in_spatial_dims=in_spatial_dims)
     input_data = self.input_data.copy_with_feature_last()
     spatial_axes = list(range(num_batch_dims, input_data.batch_ndim - 1))
+    if in_spatial_dims:
+      assert (
+        [input_data.dim_tags[a] for a in spatial_axes]
+        == [input_data.get_dim_tag_from_description(d) for d in in_spatial_dims])
     padding = padding.lower()
     if strides is None:
       strides = filter_size
@@ -5112,7 +5148,7 @@ def get_out_data_from_opts(cls, name, sources, network,
     :param DimensionTag|None out_dim:
     :param list[DimensionTag]|None out_spatial_dims:
     :param DimensionTag|None in_dim:
-    :param list[DimensionTag]|None in_spatial_dims:
+    :param list[DimensionTag|str]|None in_spatial_dims:
     :rtype: Data
     """
     input_data = get_concat_sources_data_template(sources)
diff --git a/returnn/tf/util/data.py b/returnn/tf/util/data.py
--- a/returnn/tf/util/data.py
+++ b/returnn/tf/util/data.py
@@ -2437,6 +2437,9 @@ def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None):
 
   def copy_split_feature_dim(self, new_feature_dim):
     """
+    Split it into (new_spatial_dim, new_feat_dim), in that order.
+    This will increase the feature_dim_axis by one.
+
     :param int new_feature_dim: will be the new dim
     :rtype: Data
     """
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -209,7 +209,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 7
+  _latest_behavior_version = 8
   _behavior_version = None  # type: typing.Optional[int]
 
   @classmethod


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -736,11 +736,13 @@ def test_cnn_building_block():
         "swap_axes": {
           "class": "swap_axes", "axis1": "dim:%i" % channel_num, "axis2": "dim:%i" % feature_dim, "from": "split"},
         "c1": {"class": "conv", "n_out": filters, "filter_size": filter_size, "auto_use_channel_first": False,
+               "in_spatial_dims": ("T", "dim:6"),
                "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
                "from": "swap_axes"},
         "bn1": {"class": "batch_norm", "from": "c1"},
         "y1": {"class": "activation", "activation": "relu", "batch_norm": False, "from": "bn1"},
         "c2": {"class": "conv", "n_out": filters, "filter_size": filter_size, "auto_use_channel_first": False,
+               "in_spatial_dims": ("T", "dim:6"),
                "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
                "from": "y1"},
         "p": {"class": "combine", "kind": "add", "from": ["c2", "swap_axes"]},
@@ -3567,26 +3569,26 @@ def test_conv_layer_NCHW():
         name="src_nhwc", network=net,
         output=Data(**{
           "name": "src_nhwc_output",
-          "dim": 16,
-          "shape": (None, 16, 16),
+          "dim": 17,
+          "shape": (None, 16, 17),
           "batch_dim_axis": 0,
           "time_dim_axis": 1,
           "feature_dim_axis": 3,
           "sparse": False}))
-      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 16), dtype=tf.float32)
+      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 17), dtype=tf.float32)
       src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
     with tf_compat.v1.variable_scope("src_nchw"):
       src_nchw = InternalLayer(
         name="src_nchw", network=net,
         output=Data(**{
           "name": "src_nchw_output",
-          "dim": 16,
-          "shape": (16, None, 16),
+          "dim": 17,
+          "shape": (17, None, 16),
           "batch_dim_axis": 0,
           "time_dim_axis": 2,
           "feature_dim_axis": 1,
           "sparse": False}))
-      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None, 16), dtype=tf.float32)
+      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 17, None, 16), dtype=tf.float32)
       src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
 
     filters = 64
@@ -3597,31 +3599,37 @@ def test_conv_layer_NCHW():
     with tf_compat.v1.variable_scope("conv_nhwc_from_nhwc"):
       conv_nhwc_from_nhwc = ConvLayer(
         name="conv_nhwc_from_nhwc", network=net, n_out=filters, filter_size=filter_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, auto_use_channel_first=False, sources=[src_nhwc],
         output=ConvLayer.get_out_data_from_opts(name="conv_nhwc_from_nhwc", n_out=filters,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 filter_size=filter_size, padding=padding, strides=strides,
                                                 auto_use_channel_first=False,
                                                 network=net, sources=[src_nhwc]))
     with tf_compat.v1.variable_scope("conv_nchw_from_nhwc"):
       conv_nchw_from_nhwc = ConvLayer(
         name="conv_nchw_from_nhwc", network=net, n_out=filters, filter_size=filter_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, auto_use_channel_first=True, sources=[src_nhwc],
         output=ConvLayer.get_out_data_from_opts(name="conv_nchw_from_nhwc", n_out=filters,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 filter_size=filter_size, padding=padding, strides=strides,
                                                 auto_use_channel_first=True,
                                                 network=net, sources=[src_nhwc]))
     with tf_compat.v1.variable_scope("conv_nchw_from_nchw"):
       conv_nchw_from_nchw = ConvLayer(
         name="conv_nchw_from_nchw", network=net, n_out=filters, filter_size=filter_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, auto_use_channel_first=True, sources=[src_nchw],
         output=ConvLayer.get_out_data_from_opts(name="conv_nchw_from_nchw", n_out=filters,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 filter_size=filter_size, padding=padding, strides=strides,
                                                 auto_use_channel_first=True,
                                                 network=net, sources=[src_nchw]))
     tf_compat.v1.global_variables_initializer().run()
     out, seq_lens = session.run([conv_nhwc_from_nhwc.output.placeholder,
                                  conv_nhwc_from_nhwc.output.size_placeholder[0]],
-                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 16),
+                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 17),
                                            src_nhwc.output.size_placeholder[0]: np.full(shape=(10,), fill_value=10)}
                                 )
     print(out.shape)
@@ -3630,7 +3638,7 @@ def test_conv_layer_NCHW():
     time_dim_axis = 1 if tf_util.is_gpu_available() else 0
     out, seq_lens = session.run([conv_nchw_from_nhwc.output.placeholder,
                                  conv_nchw_from_nhwc.output.size_placeholder[time_dim_axis]],
-                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 16),
+                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 17),
                                            src_nhwc.output.size_placeholder[0]: np.full(shape=(10,), fill_value=10)
                                 })
     print(out.shape)
@@ -3642,7 +3650,7 @@ def test_conv_layer_NCHW():
     if tf_util.is_gpu_available():
       out, seq_lens = session.run([conv_nchw_from_nchw.output.placeholder,
                                    conv_nchw_from_nchw.output.size_placeholder[1]],
-                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 16, 10, 16),
+                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 17, 10, 16),
                                              src_nchw.output.size_placeholder[1]: np.full(shape=(10,), fill_value=10)
                                   })
       print(out.shape)
@@ -3659,26 +3667,26 @@ def test_pool_layer_NCHW():
         name="src_nhwc", network=net,
         output=Data(**{
           "name": "src_nhwc_output",
-          "dim": 16,
-          "shape": (None, 16, 16),
+          "dim": 17,
+          "shape": (None, 16, 17),
           "batch_dim_axis": 0,
           "time_dim_axis": 1,
           "feature_dim_axis": 3,
           "sparse": False}))
-      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 16), dtype=tf.float32)
+      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 17), dtype=tf.float32)
       src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
     with tf_compat.v1.variable_scope("src_nchw"):
       src_nchw = InternalLayer(
         name="src_nchw", network=net,
         output=Data(**{
           "name": "src_nchw_output",
-          "dim": 16,
-          "shape": (16, None, 16),
+          "dim": 17,
+          "shape": (17, None, 16),
           "batch_dim_axis": 0,
           "time_dim_axis": 2,
           "feature_dim_axis": 1,
           "sparse": False}))
-      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None, 16), dtype=tf.float32)
+      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 17, None, 16), dtype=tf.float32)
       src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
 
     pool_size = (5, 5)
@@ -3688,55 +3696,61 @@ def test_pool_layer_NCHW():
     with tf_compat.v1.variable_scope("pool_nhwc_from_nhwc"):
       pool_nhwc_from_nhwc = PoolLayer(
         name="pool_nhwc_from_nhwc", network=net, mode="max", pool_size=pool_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, use_channel_first=False, sources=[src_nhwc],
         output=PoolLayer.get_out_data_from_opts(name="pool_nhwc_from_nhwc",
                                                 pool_size=pool_size, padding=padding, strides=strides,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 use_channel_first=False,
                                                 network=net, sources=[src_nhwc]))
     with tf_compat.v1.variable_scope("pool_nchw_from_nhwc"):
       pool_nchw_from_nhwc = PoolLayer(
         name="pool_nchw_from_nhwc", network=net, mode="max", pool_size=pool_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, use_channel_first=True, sources=[src_nhwc],
         output=PoolLayer.get_out_data_from_opts(name="pool_nchw_from_nhwc",
                                                 pool_size=pool_size, padding=padding, strides=strides,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 use_channel_first=True,
                                                 network=net, sources=[src_nhwc]))
     with tf_compat.v1.variable_scope("pool_nchw_from_nchw"):
       pool_nchw_from_nchw = PoolLayer(
         name="pool_nchw_from_nchw", network=net, mode="max", pool_size=pool_size,
+        in_spatial_dims=["T", "dim:16"],
         padding=padding, strides=strides, use_channel_first=True, sources=[src_nchw],
         output=PoolLayer.get_out_data_from_opts(name="pool_nchw_from_nchw",
                                                 pool_size=pool_size, padding=padding, strides=strides,
+                                                in_spatial_dims=["T", "dim:16"],
                                                 use_channel_first=True,
                                                 network=net, sources=[src_nchw]))
     tf_compat.v1.global_variables_initializer().run()
     out, seq_lens = session.run([pool_nhwc_from_nhwc.output.placeholder,
                                  pool_nhwc_from_nhwc.output.get_sequence_lengths()],
-                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 16),
+                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 17),
                                            src_nhwc.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)}
                                 )
     print(out.shape)
-    assert_equal(out.shape, (10, 7, 6, 16))
+    assert_equal(out.shape, (10, 7, 6, 17))
     print(seq_lens)
     out, seq_lens = session.run([pool_nchw_from_nhwc.output.placeholder,
                                  pool_nchw_from_nhwc.output.get_sequence_lengths()],
-                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 16),
+                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 17),
                                            src_nhwc.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)
                                 })
     print(pool_nchw_from_nhwc.output, out.shape)
     if pool_nchw_from_nhwc.output.feature_dim_axis == 1:
-      assert_equal(out.shape, (10, 16, 7, 6))
+      assert_equal(out.shape, (10, 17, 7, 6))
     else:
-      assert_equal(out.shape, (10, 7, 6, 16))
+      assert_equal(out.shape, (10, 7, 6, 17))
     print(seq_lens)
     if tf_util.is_gpu_available():
       out, seq_lens = session.run([pool_nchw_from_nchw.output.placeholder,
                                    pool_nchw_from_nchw.output.get_sequence_lengths()],
-                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 16, 11, 16),
+                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 17, 11, 16),
                                              src_nchw.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)
                                   })
       print(out.shape)
-      assert_equal(out.shape, (10, 16, 7, 6))
+      assert_equal(out.shape, (10, 17, 7, 6))
       print(seq_lens)
 
 
@@ -3756,6 +3770,7 @@ def test_TransposedConvLayer_2d_simple():
       "output": {  # [B,D',T,2]
         'class': 'transposed_conv', 'from': 'Unflatten',
         'activation': None, 'with_bias': True,
+        "in_spatial_dims": ("T", "dim:1"),
         'n_out': n_out, 'filter_size': (1, 2), 'strides': (1, 1),
         'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
     })
@@ -3785,6 +3800,7 @@ def test_TransposedConvLayer_2d_2x2():
       "output": {  # [B,D',T,2]
         'class': 'transposed_conv', 'from': 'Unflatten',
         'activation': None, 'with_bias': True,
+        "in_spatial_dims": ("T", "dim:1"),
         'n_out': n_out, 'filter_size': (2, 2), 'strides': (2, 2),
         'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
     })