SplitDimsLayer fix feature_dim_axis on feature-dim split

[albertz]

Fix #704.

WARNING: This potentially changed the behavior of configs.
E.g. the one reported in #703.
Most of the attention and transducer encoders which used initial convolutional layers.

Related is #596. Maybe actually the fix for that in #683 introduces this new bug.

[albertz]

So, Git bisecting:

It is bad in current master (ea3e91f).
It was still good in 6f7f622 (10 days ago).
2ff056d is the first bad commit.

commit 2ff056de31838b6425333a882759df0093b5bb63
Author: Albert Zeyer <[email protected]>
Date:   Wed Sep 22 14:45:18 2021 +0200

    SplitDimsLayer, fix feature_dim_axis for batch-feature-major input
    
    Fix #596

 returnn/tf/layers/basic.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

So yes, as expected #683 introduced this bug.

[albertz]

Note on the potential wrong behavior

This did not cause an error, just wrong behavior!
When you used split dims like this:

"source0": {"class": "split_dims", "axis": "F", "dims": (-1, 1), "from": "data"},  # (B,T,40,1)

And then 2D convolution, it would operate on the two expected spatial dims. The spatial dims are those which are not the feature dim and the batch dim. So it matters, from these splitted dims 40 -> (40,1), which of these are actually the feature dim.

Earlier, the last axis (with dim 1) was the new feature dim.
#683 changed this, and dim 40 here kept being the feature dim, as you can see in this output:

layer root/'source0' output: Data{'source0_output', [B,T|'time'[B],F|F'feature:data'(40),'source0_split_dims1'(1)]}

In some follow up 2D conv with options:

"conv0": {
  "class": "conv", "from": "source0",
  "padding": "same", "filter_size": (3, 3),
  "n_out": 32, "activation": None, "with_bias": True},  # (B,T,40,32)

you then see:

layer root/'conv0' output: Data{'conv0_output', [B,T|'time'[B],'source0_split_dims1'(1),F|F'conv0:channel'(32)]}

So the second spatial dim this operated on was source0_split_dims1 from above.
This was not the intended and old behavior.

[albertz]

Note that with #706 (out_shape for verification) we could have detected such problem more directly.

And the proposed options in #597 (e.g. in_spatial_dims for ConvLayer) also would have avoided the problem.

Probably any solution for #594 would also have avoided the problem.