Repository: rwth-i6/returnn

Issue:
RF TF-layers: `Dim.complete_dyn_size`, `Tensor.copy_transpose` etc: order of dims ambiguous
In the RETURNN frontend (RF), the TF layers dict backend is kind of special in that the raw tensor (the layer dict) does not have a well defined order of dimensions, so it means that many raw ops like `Backend.get_shape_tuple_raw` don't make sense, and neither ops like `Tensor.copy_transpose`, `Tensor.copy_compatible_to`, `Tensor.copy_compatible_to_dims` etc. All such ops in the backend raise some exception (currently `NotImplementedError`, although maybe some other exception might make more sense).

There is one exception, which is `Backend.transpose`, which uses the `TransposeLayer`, which makes the order of dimensions well defined explicitly. This is used for the final output, but it really makes sense only there, nowhere else, and we specifically don't want to have this anywhere else.

Usually, we should also not ever get to such  op/function being used. However, currently the logic in `Dim.complete_dyn_size` does use `copy_compatible_to_dims`. Currently, the code works fine anyway because in case of a scalar of in case of just [B] shape, the order of dims is not relevant, and there is a check in `copy_compatible_to_dims` which would just return `copy` if the order would not change, so nothing happens. So that's why all tests are passing, and probably most use cases. However, for more dims, the code is already broken. Now with potential rewrites of `copy_compatible_to_dims` (#1409), this becomes more relevant.

The question is, how to deal with this in general.

- `Backend.transpose` and `Tensor.copy_transpose` seem to be the same, but `Backend.transpose` has this special usage for the final output, while `Tensor.copy_transpose` usage is spread widely across RETURNN. Maybe we should rename `Backend.transpose` to `Backend.make_output_tensor` or so to make this use case explicitly clear.
- `Tensor.copy_transpose` etc, should they just ignore the order and return `self.copy()`, in case we have a backend where the order of dims in the raw tensor is undefined? So we need a new function in the backend, like `Backend.has_defined_raw_order_of_dims()` or so.
- `Data.get_common_data` and then `Tensor.copy_compatible_to` or `Tensor.copy_compatible_to_dims`, or `bin_op_out_template`, how should they deal with adding broadcast dimensions? Also just ignore? Ignoring them would be fine only if we don't need the flag `allow_broadcast_all_sources`.



Discussion:
For reference, an example stack trace from such error ([via](https://github.com/rwth-i6/returnn/actions/runs/6431758167/job/17465274219), #1409):

```
  File "/home/runner/work/returnn/returnn/tests/test_rf_conv.py", line 299, in test_maxpool1d_stride_border_cond.<locals>._Net.__call__
    line: return rf.max_pool1d(x, pool_size=6, strides=3, padding="valid", in_spatial_dim=in_spatial_dim)
    locals:
      rf = <global> <module 'returnn.frontend' from '/home/runner/work/returnn/returnn/returnn/frontend/__init__.py'>
      rf.max_pool1d = <global> <function max_pool1d at 0x7f6f16b467a0>
      x = <local> Tensor{'data', [B,'time'[B],'in'(7)]}
      pool_size = <not found>
      strides = <not found>
      padding = <not found>
      in_spatial_dim = <local> Dim{'time'[B]}
  File "/home/runner/work/returnn/returnn/returnn/frontend/conv.py", line 498, in max_pool1d
    line: out_spatial_dim=out_spatial_dim,
    locals:
      out_spatial_dim = <local> None
  File "/home/runner/work/returnn/returnn/returnn/frontend/conv.py", line 535, in pool1d
    line: out_spatial_dims=[out_spatial_dim] if out_spatial_dim is not None else None,
    locals:
      out_spatial_dims = <not found>
      out_spatial_dim = <local> None
  File "/home/runner/work/returnn/returnn/returnn/frontend/conv.py", line 451, in pool
    line: out_spatial_dims=out_spatial_dims,
    locals:
      out_spatial_dims = <local> None
  File "/home/runner/work/returnn/returnn/returnn/tf/frontend_layers/_backend.py", line 1021, in ReturnnLayersBackend.pool
    line: padding=padding,
    locals:
      padding = <local> 'valid'
  File "/home/runner/work/returnn/returnn/returnn/frontend/conv.py", line 646, in make_conv_out_spatial_dims
    line: padding=padding,
    locals:
      padding = <local> 'valid'
  File "/home/runner/work/returnn/returnn/returnn/frontend/conv.py", line 698, in _calc_out_dim
    line: valid_part = in_dim.sub_left(filter_left_dilated).sub_right(filter_right_dilated)
    locals:
      valid_part = <not found>
      in_dim = <local> Dim{'time'[B]}
      in_dim.sub_left = <local> <bound method _DimMixin.sub_left of Dim{'time'[B]}>
      filter_left_dilated = <local> 2
      sub_right = <not found>
      filter_right_dilated = <local> 3
  File "/home/runner/work/returnn/returnn/returnn/tensor/_dim_extra.py", line 1952, in _DimMixin.sub_left
    line: return term.as_dim()
    locals:
      term = <local> Dim._OpLinearTerm([Dim._OpMultTerm([Dim{'unnamed_dim_-2'(-2)}]), Dim._OpMultTerm([Dim{'time'[B]}])])
      term.as_dim = <local> <bound method _OpLinearTerm.as_dim of Dim._OpLinearTerm([Dim._OpMultTerm([Dim{'unnamed_dim_-2'(-2)}]), Dim._OpMultTerm([Dim{'time'[B]}])])>
  File "/home/runner/work/returnn/returnn/returnn/tensor/_dim_extra.py", line 2509, in _OpLinearTerm.as_dim
    line: derived_from_tag=self.representative_tag(),
    locals:
      derived_from_tag = <not found>
      self = <local> Dim._OpLinearTerm([Dim._OpMultTerm([Dim{'unnamed_dim_-2'(-2)}]), Dim._OpMultTerm([Dim{'time'[B]}])])
      self.representative_tag = <local> <bound method _OpLinearTerm.representative_tag of Dim._OpLinearTerm([Dim._OpMultTerm([Dim{'unnamed_dim_-2'(-2)}]), Dim._OpMultTerm([Dim{'time'[B]}])])>
  File "/home/runner/work/returnn/returnn/returnn/tensor/dim.py", line 90, in Dim.__init__
    line: self._handle_extra_kwargs(**kwargs)
    locals:
      self = <local> Dim{'-2+time'[?]}
      self._handle_extra_kwargs = <local> <bound method _DimMixin._handle_extra_kwargs of Dim{'-2+time'[?]}>
      kwargs = <local> {'kind': <spatial>, 'derived_from_op': <Dim.Op 'add' [Dim{'unnamed_dim_-2'(-2)}, Dim{'time'[B]}]>, 'derived_from_tag': Dim{'time'[B]}}
  File "/home/runner/work/returnn/returnn/returnn/tensor/_dim_extra.py", line 146, in _DimMixin._handle_extra_kwargs
    line: self.complete_dyn_size()
    locals:
      self = <local> Dim{'-2+time'[?]}
      self.complete_dyn_size = <local> <bound method _DimMixin.complete_dyn_size of Dim{'-2+time'[?]}>
  File "/home/runner/work/returnn/returnn/returnn/tensor/_dim_extra.py", line 1153, in _DimMixin.complete_dyn_size
    line: x_ = x.copy_compatible_to_dims(common.dims) if x.dims else x
    locals:
      x_ = <not found>
      x = <local> Tensor{'time', [B], dtype='int32'}
      x.copy_compatible_to_dims = <local> <bound method tensor_copy_compatible_to_dims of Tensor{'time', [B], dtype='int32'}>
      common = <local> Tensor{'time', [B], dtype='int32'}
      common.dims = <local> (Dim{B},)
      x.dims = <local> (Dim{B},)
  File "/home/runner/work/returnn/returnn/returnn/tf/frontend_layers/_backend.py", line 153, in ReturnnLayersBackend.get_shape_tuple_raw
    line: raise NotImplementedError
    locals:
      NotImplementedError = <builtin> <class 'NotImplementedError'>
NotImplementedError
```

Resolving patch:
diff --git a/returnn/tensor/_dim_extra.py b/returnn/tensor/_dim_extra.py
--- a/returnn/tensor/_dim_extra.py
+++ b/returnn/tensor/_dim_extra.py
@@ -1089,18 +1089,17 @@ def _bin_op(a, b):
                     return a
                 if isinstance(b, _t.Tensor):
                     return b
+                raise Exception(f"Dim complete_dyn_size: bin_op: expect to get one Tensor, got {a} and {b}")
             if kind == "add":
-                return _relu(a + b)
+                return _relu(rf.combine_bc(a, "add", b))
             elif kind == "sub":
-                return _relu(a - b)
+                return _relu(rf.combine_bc(a, "sub", b))
             elif kind == "mul":
-                return a * b
+                return rf.combine_bc(a, "mul", b)
             elif kind in ("floordiv", "truediv"):  # truediv assumes there is no remainder
-                return a // b
+                return rf.combine_bc(a, "floordiv", b)
             elif kind == "ceildiv":
-                if isinstance(a, _t.Tensor):
-                    return rf.ceil_divide(a, b)
-                return -(-a // b)
+                return -rf.combine_bc(-a, "floordiv", b)
             else:
                 raise ValueError("unknown op kind %r" % op.kind)
 
@@ -1148,17 +1147,14 @@ def _relu(a):
             if y is None:
                 y = x.copy(name=y_name)
                 continue
-            if x.dim_tags != y.dim_tags:
+            if tf:
                 common = _t.Tensor.get_common_data([x, y], allow_broadcast_all_sources=True)
                 x_ = x.copy_compatible_to_dims(common.dims) if x.dims else x
                 y_ = y.copy_compatible_to_dims(common.dims) if y.dims else y
                 y = common
-            else:
-                x_, y_ = x, y
-            if tf:
                 y.placeholder = _bin_op_tf(y_.placeholder, x_.placeholder)
             else:
-                y = _bin_op(y_, x_)
+                y = _bin_op(y, x)
         assert y, f"op {op}?"
         if self.dyn_size_ext:
             assert self.dyn_size_ext.dim_tags == y.dim_tags


Test patch: