Repository: rwth-i6/returnn

Issue:
`SplitDimsLayer` wrong out feature axis when splitting feature axis
`data` has shape [B,T,40]

`SplitDimsLayer` config:
```
"source0": {"class": "split_dims", "axis": "F", "dims": (-1, 1), "from": "data"},  # (T,40,1)
```

Results in:
```
layer root/'source0' output: Data{'source0_output', [B,T|'time'[B],F|F'feature:data'(40),'source0_split_dims1'(1)]}
```

Earlier, the feature dim axis was the last dim (i.e. the `source0_split_dims1` here).

Related is #596. Maybe actually the fix for that in #683 introduces this new bug.

(I found this bug in #703 but this is probably otherwise unrelated.)



Discussion:

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
@@ -3058,11 +3058,12 @@ 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, split_offset=None):
+  def _map_old_axis_to_new_axis(cls, split_axis, dims, old_axis, use_remaining=True, split_offset=None):
     """
     :param int split_axis:
     :param tuple[int] dims: might include -1
     :param int old_axis:
+    :param bool use_remaining: whether to make use of -1 in dims
     :param int|None split_offset:
     :rtype: int
     """
@@ -3071,7 +3072,7 @@ def _map_old_axis_to_new_axis(cls, split_axis, dims, old_axis, split_offset=None
     if old_axis > split_axis:
       return old_axis + len(dims) - 1
     assert old_axis == split_axis
-    if -1 in dims:
+    if use_remaining and -1 in dims:
       assert dims.count(-1) == 1
       return split_axis + dims.index(-1)
     assert split_offset is not None
@@ -3127,9 +3128,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
+    else:
+      out.time_dim_axis = cls._map_old_axis_to_new_axis(
+        split_axis=axis, dims=dims, old_axis=data.time_dim_axis, use_remaining=True, split_offset=0)
     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)
+        split_axis=axis, dims=dims, old_axis=data.feature_dim_axis, use_remaining=False, 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]


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -1566,13 +1566,14 @@ def test_SplitDimsLayer_batch_feature_major_keep_feature():
     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]}
+      "output": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]}  # [B,D,T,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
+    assert out.time_dim_axis == 2
     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)
@@ -1580,6 +1581,28 @@ def test_SplitDimsLayer_batch_feature_major_keep_feature():
     numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))
 
 
+def test_SplitDimsLayer_split_feature():
+  n_batch, n_time, n_in = 7, 3, 5
+  config = Config({
+    "extern_data": {
+      "data": {"dim": n_in}  # [B,T,D]
+    }
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict({
+      "output": {'class': 'split_dims', 'from': 'data', 'axis': 'F', 'dims': [-1, 1]}
+    })
+    out = net.get_default_output_layer().output
+    assert out.batch_shape == (None, None, n_in, 1)
+    assert out.feature_dim_axis == 3  # https://github.com/rwth-i6/returnn/issues/704
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
+    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_time, n_in, 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):
   """