Repository: rwth-i6/returnn

Issue:
`SplitDimsLayer` wrong feature axis
Via https://github.com/rwth-i6/pytorch-to-returnn/pull/54.

```
# data: [B,F|F'feature:data'(512),T|'time:var:extern_data:data'[B]]
"Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},
# unflatten: [B,F'feature:data'(512),T|'time:var:extern_data:data'[B],F|'Unflatten_split_dims1'(1)]
```

I would expect that the original feature dim axis stays as before.

This is due to the default behavior of `NotSpecified`.
As it is not specified (`NotSpecified`) before, it stays not specified after.
But that changes the feature dim axis (1 before, then -1 after).



Discussion:
@Zettelkasten what do you think? This should be changed, right?
Ouh yes, this is against all principles that the order of axes should not matter.
In this case, maybe even in general, I am also quite sure that nobody relied on this behavior.
From my POV, passing `feature_dim_axis=NotSpecified` to the `Data` should be resolved to some concrete `int|None`, just like it is now for `time_dim_axis`. But that leaves some questions open, e.g. what should happen if you take a `[B,T]` tensor with `feature_dim_axis=NotSpecified`, and than expand dims by 1. Should the new dim be the feature dim or not? And so on ...
Very related is #586.

But then again I avoided using `NotSpecified` in all my configs, I never found it useful, only non-transparent. Maybe people that use this have other opinions, and use/like this "feature"?


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
@@ -3054,23 +3054,27 @@ def _resolve_dims(cls, old_dim, new_dims, pad_to_multiples=False):
     return tuple(new_dims_resolved)
 
   @classmethod
-  def _map_old_axis_to_new_axis(cls, split_axis, dims, old_axis, require_exact=True):
+  def _map_old_axis_to_new_axis(cls, split_axis, dims, old_axis, split_offset=None):
     """
     :param int split_axis:
     :param tuple[int] dims: might include -1
     :param int old_axis:
-    :param bool require_exact: if True, must be unique, always returns int, or raise exception
+    :param int|None split_offset:
     :rtype: int
     """
     if old_axis < split_axis:
       return old_axis
     if old_axis > split_axis:
       return old_axis + len(dims) - 1
+    assert old_axis == split_axis
     if -1 in dims:
       assert dims.count(-1) == 1
-      return old_axis + dims.index(-1)
-    assert not require_exact
-    return old_axis
+      return split_axis + dims.index(-1)
+    assert split_offset is not None
+    if split_offset < 0:
+      split_offset += len(dims)
+    assert 0 <= split_offset < len(dims)
+    return split_axis + split_offset
 
   @classmethod
   def get_out_data_from_opts(cls, name, axis, dims, pad_to_multiples=None, sources=(), **kwargs):
@@ -3119,6 +3123,12 @@ def get_out_data_from_opts(cls, name, axis, dims, pad_to_multiples=None, sources
     out = data.copy_template_new_dim_tags(new_dim_tags)
     if data.time_dim_axis is None:
       out.time_dim_axis = None
+    if data.feature_dim_axis is not None:
+      expected_out_feature_dim_axis = cls._map_old_axis_to_new_axis(
+        split_axis=axis, dims=dims, old_axis=data.feature_dim_axis, split_offset=-1)
+      if out.feature_dim_axis != expected_out_feature_dim_axis:  # maybe due to non-specified default behavior
+        out.feature_dim_axis = expected_out_feature_dim_axis
+        out.dim = out.batch_shape[out.feature_dim_axis]
     return out
 
 


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -1533,6 +1533,31 @@ def test_SplitDimsLayer_resolve_dims():
   assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5 + 2, new_dims=(2, 3, -1, 1), pad_to_multiples=True), (2, 3, 6, 1))
 
 
+def test_SplitDimsLayer_batch_feature_major_keep_feature():
+  n_batch, n_time, n_in = 7, 3, 5
+  config = Config({
+    "extern_data": {
+      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
+    }
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    assert net.extern_data.get_default_input_data().is_batch_feature_major
+    net.construct_from_dict({
+      "output": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]}
+    })
+    out = net.get_default_output_layer().output
+    assert_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
+    assert out.dim_tags[1].dimension == n_in and out.dim_tags[3].dimension == 1
+    assert out.placeholder.shape.as_list() == [None, n_in, None, 1]
+    assert out.feature_dim_axis == 1  # https://github.com/rwth-i6/returnn/issues/596
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
+    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
+    assert isinstance(out_v, numpy.ndarray)
+    assert out_v.shape == (n_batch, n_in, n_time, 1)
+    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))
+
+
 def _check_MergeDimsLayer(session, in_data_opts, in_static_shape, opts, out_data_shape, out_static_shape,
                           in_sizes=None, out_sizes=None):
   """