Repository: rwth-i6/returnn

Issue:
RETURNN frontend: no module named tensorflow, but backend is torch
Good morning!

I was testing the latest RETURNN frontend `rf`, and I stumbled upon the following issue: `ModuleNotFoundError: No module named 'tensorflow'`. However, I'm running the `demo-rf-config` with `backend = "torch"`. In an environment with `tensorflow`, it completely loads TF so the startup is slower, but the loss converges to the same number.

Complete stack trace (omitting `<frame not available>` as `...`):

```
<module>, __init__.py:4
<frame not available>
...
<frame not available>
copy_add_dim_by_tag, _tensor_extra.py:912
copy_compatible_to, _tensor_extra.py:1116
bin_op_out_template, _utils.py:122
combine, _backend.py:525
combine, math_.py:182
__add__, _tensor_op_overloads.py:40
__call__, linear.py:38
__call__, demo-rf.config:41
train_step, demo-rf.config:51
_run_step, engine.py:292
train_epoch, engine.py:147
train, engine.py:119
execute_main_task, __main__.py:419
main, __main__.py:595
<module>, rnn.py:11
```

The `copy_add_dim_by_tag` function seems to be running `from returnn.tf.util.data import BatchInfo` at the line indicated by the stack trace (`returnn/frontend/_tensor_extra.py:912`).

Python version: 3.10.6
PyTorch version: 1.13.1

If there's any other detail I can give, please let me know.


Discussion:

Resolving patch:
diff --git a/returnn/tensor/_tensor_extra.py b/returnn/tensor/_tensor_extra.py
--- a/returnn/tensor/_tensor_extra.py
+++ b/returnn/tensor/_tensor_extra.py
@@ -787,18 +787,12 @@ def copy_with_feature_last(self) -> _t.Tensor:
 
     def copy_add_batch_dim(self, batch_dim_axis, batch=None, dim_tag=None) -> _t.Tensor:
         """
-        Warning: Assumes TensorFlow.
-
         :param int batch_dim_axis:
         :param BatchInfo|None batch:
         :param Dim|None dim_tag:
         :return: copy of myself with added batch-dim
         """
         assert self.batch_dim_axis is None
-        if not batch:
-            from returnn.tf.layers.base import LayerBase
-
-            batch = LayerBase.get_recent_layer().get_batch_info()
         if batch_dim_axis < 0:
             assert batch_dim_axis + self.batch_ndim + 1 >= 0
             batch_dim_axis += self.batch_ndim + 1
@@ -808,19 +802,36 @@ def copy_add_batch_dim(self, batch_dim_axis, batch=None, dim_tag=None) -> _t.Ten
         if placeholder is not None:
             backend = self._raw_backend
             placeholder = backend.expand_dims_raw(placeholder, batch_dim_axis)
-            if not isinstance(batch.dim, int) or batch.dim != 1:
-                placeholder = backend.expand_raw(placeholder, batch_dim_axis, batch.dim)
+            if batch:
+                batch_dim_ = batch.dim
+            elif dim_tag:
+                if dim_tag.dyn_size_ext:
+                    assert dim_tag.dyn_size_ext.dims == ()
+                    assert dim_tag.dyn_size_ext.raw_tensor is not None
+                    batch_dim_ = dim_tag.dyn_size_ext.raw_tensor
+                elif dim_tag.dimension:
+                    batch_dim_ = dim_tag.dimension
+                else:
+                    raise Exception(f"{self} copy_add_batch_dim: unknown batch dim for {dim_tag!r}")
+            else:
+                raise Exception(f"{self} copy_add_batch_dim: unknown batch dim ")
+            if not isinstance(batch_dim_, int) or batch_dim_ != 1:
+                placeholder = backend.expand_raw(placeholder, batch_dim_axis, batch_dim_)
         dim_tags = list(self.dim_tags)
         if dim_tag:
             assert dim_tag.is_batch_dim()
-            assert dim_tag.dimension == batch.static_dim or dim_tag.dimension is None
             assert dim_tag.batch == batch
+            if batch:
+                assert dim_tag.dimension == batch.static_dim or dim_tag.dimension is None
         else:
-            dim_tag = Dim(kind=Dim.Types.Batch, description="batch", dimension=batch.static_dim, batch=batch)
+            dim_tag = Dim(
+                kind=Dim.Types.Batch, description="batch", dimension=batch.static_dim if batch else None, batch=batch
+            )
         dim_tags.insert(batch_dim_axis, dim_tag)
         data_opts["dims"] = dim_tags
-        data_opts["batch"] = batch
-        data_opts["beam"] = batch.beam
+        if batch:
+            data_opts["batch"] = batch
+            data_opts["beam"] = batch.beam
         other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
         for k, a in other_special_axes.items():
             data_opts[k] = a if (a < batch_dim_axis) else (a + 1)
@@ -909,14 +920,17 @@ def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None) -> _t.Tenso
             if unbroadcast:
                 return self.copy_add_batch_dim(batch_dim_axis=axis, batch=dim_tag.batch, dim_tag=dim_tag)
             else:
-                from returnn.tf.util.data import BatchInfo
+                if dim_tag.batch or self.batch:
+                    from returnn.tf.util.data import BatchInfo
 
-                batch_info = BatchInfo.make_global_broadcast_batch_info()
-                return self.copy_add_batch_dim(
-                    batch_dim_axis=axis,
-                    batch=batch_info,
-                    dim_tag=dim_tag if (dim_tag.dimension == 1 and dim_tag.batch == batch_info) else None,
-                )
+                    batch_info = BatchInfo.make_global_broadcast_batch_info()
+                else:
+                    batch_info = None
+                if dim_tag and dim_tag.dimension == 1 and dim_tag.batch == batch_info:
+                    pass  # keep it
+                else:
+                    dim_tag = Dim(kind=Dim.Types.Batch, description="batch-broadcast", dimension=1, batch=batch_info)
+                return self.copy_add_batch_dim(batch_dim_axis=axis, batch=batch_info, dim_tag=dim_tag)
 
         data_opts = self.get_kwargs()
         # Note: if dim_tag is feature, but we are sparse, we just make it spatial
@@ -2812,9 +2826,11 @@ def get_static_batch_dim(self):
         :rtype: int|None
         """
         # Do not fallback to get_batch_dim or get_recent_layer or so. This should be safe.
-        if not self.batch:
-            return None
-        return self.batch.static_dim
+        if self.batch:
+            return self.batch.static_dim
+        if self.have_batch_axis():
+            return self.get_batch_dim_tag().dimension
+        return None
 
     def get_spatial_batch_axes(self):
         """


Test patch:
diff --git a/tests/_setup_test_env.py b/tests/_setup_test_env.py
--- a/tests/_setup_test_env.py
+++ b/tests/_setup_test_env.py
@@ -64,13 +64,16 @@ def _chained_excepthook(exctype, value, traceback):
     log.initialize(verbosity=[5], propagate=False)
 
     # TF is optional.
-    # Note that importing TF still has a small side effect:
-    # BackendEngine._get_default_engine() will return TF by default, if TF is already loaded.
-    # For most tests, this does not matter.
-    try:
-        import tensorflow as tf
-    except ImportError:
+    if "RETURNN_DISABLE_TF" in os.environ and int(os.environ["RETURNN_DISABLE_TF"]) == 1:
         tf = None
+    else:
+        # Note that importing TF still has a small side effect:
+        # BackendEngine._get_default_engine() will return TF by default, if TF is already loaded.
+        # For most tests, this does not matter.
+        try:
+            import tensorflow as tf
+        except ImportError:
+            tf = None
 
     if tf:
         import returnn.tf.util.basic as tf_util
diff --git a/tests/test_TFUtil.py b/tests/test_TFUtil.py
--- a/tests/test_TFUtil.py
+++ b/tests/test_TFUtil.py
@@ -424,7 +424,7 @@ def test_Data_find_matching_dim_map_broadcast_matches():
     assert mapping[0] == 1 and mapping[1] == 2
 
     copied = d2.copy_compatible_to(d1)
-    assert copied.batch_ndim == d1.batch_ndim and copied.batch.static_dim == 1 and copied.batch_shape == (1, 5, 1)
+    assert copied.batch_ndim == d1.batch_ndim and copied.get_static_batch_dim() == 1 and copied.batch_shape == (1, 5, 1)
     print("copied compatible:", copied)
 
 
@@ -633,14 +633,14 @@ def test_Data_copy_compatible_to_src_no_batch():
     d2 = Data(name="d2", shape=(), batch_dim_axis=None, time_dim_axis=None)
     d2.placeholder = tf.zeros([d if (d is not None) else 1 for d in d2.batch_shape])
     d3 = d2.copy_compatible_to(d1)
-    assert d3.batch.static_dim == 1 and d3.batch_shape == (1, 1, 1)
+    assert d3.get_static_batch_dim() == 1 and d3.batch_shape == (1, 1, 1)
 
 
 def test_Data_copy_compatible_to_add_batch_dim():
     common_data = Data(name="accum_att_weights_output", shape=(None, 1))
     d1 = Data(name="att_weights_avg_output", shape=(1,), batch_dim_axis=None)
     d2 = d1.copy_compatible_to(common_data)
-    assert d2.have_batch_axis() and d2.batch.static_dim == 1 and d2.batch_shape == (1, 1, 1)
+    assert d2.have_batch_axis() and d2.get_static_batch_dim() == 1 and d2.batch_shape == (1, 1, 1)
 
 
 def test_Data_copy_compatible_to_add_feature_dim():
@@ -658,7 +658,7 @@ def test_Data_copy_compatible_to_add_batch_feature_dim():
     d2 = d1.copy_compatible_to(common_data)
     print(d2)
     assert d2.have_batch_axis() and d2.have_time_axis()
-    assert d2.batch.static_dim == 1  # batch-dim is broadcasted
+    assert d2.get_static_batch_dim() == 1  # batch-dim is broadcasted
     assert d2.batch_shape == (1, None, 1)