returnn 1.20250225.201207__tar.gz → 1.20250226.115259__tar.gz

{returnn-1.20250225.201207/returnn.egg-info → returnn-1.20250226.115259}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250225.201207
+Version: 1.20250226.115259
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250226.115259/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250226.115259'
+long_version = '1.20250226.115259+git.0d32534'

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/returnn/frontend/_backend.py

@@ -391,17 +391,17 @@ class Backend(Generic[T]):
         source: Tensor,
         *,
         dims: Sequence[Dim],
-        out_dim: Optional[Dim] = None,
-    ) -> Tuple[Tensor, Dim]:
+        out_dim: Dim,
+    ) -> Tensor:
         """
         Merges a list of axes into a single one. (Flatten the dims.)
         E.g. input is (batch, width, height, dim) and dims=(width,height), then we get (batch, width*height, dim).
         Or input is (batch, time, height, dim) and axes=(height,dim), then we get (batch, time, height*dim).
 
         :param source:
-        :param dims:
-        :param out_dim:
-        :return: tensor, out_dim
+        :param dims: list of dims to merge. len(dims) >= 2
+        :param out_dim: resulting merged dim
+        :return: tensor
         """
         raise NotImplementedError
 

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/returnn/frontend/array_.py

@@ -182,8 +182,29 @@ def merge_dims(
         else:
             out_dim = Dim(1, name="ext")
         return rf.expand_dim(source, out_dim), out_dim
+    if len(dims) == 1:
+        if out_dim is None or out_dim == dims[0]:
+            return source, dims[0]
+        return rf.replace_dim(source, in_dim=dims[0], out_dim=out_dim)
+    if out_dim is None:
+        out_dim = dims[0]
+        reset_dyn_size = False
+        for d in dims[1:]:
+            reset_dyn_size |= d.need_masking() and out_dim.capacity != 1
+            out_dim = out_dim * d
+        if reset_dyn_size:
+            # The dynamic sizes as calculated via dim math would not correctly describe how the tensor looks like.
+            # This would then potentially discard some of the data in the tensor in subsequent operations,
+            # when masking is applied.
+            # Thus, discard the dynamic sizes, and just treat it as a flat dim with scalar dynamic size.
+            # https://github.com/rwth-i6/returnn/issues/1694
+            out_dim_size = dims[0].get_dim_value_tensor()
+            for d in dims[1:]:
+                out_dim_size *= d.get_dim_value_tensor()
+            assert isinstance(out_dim_size, Tensor) and out_dim_size.dims == ()  # scalar
+            out_dim.dyn_size_ext = out_dim_size
     # noinspection PyProtectedMember
-    return source._raw_backend.merge_dims(source, dims=dims, out_dim=out_dim)
+    return source._raw_backend.merge_dims(source, dims=dims, out_dim=out_dim), out_dim
 
 
 def split_dims(

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/returnn/tf/frontend_layers/_backend.py

@@ -241,8 +241,8 @@ class ReturnnLayersBackend(Backend[Layer]):
         source: Tensor,
         *,
         dims: Sequence[Dim],
-        out_dim: Optional[Dim] = None,
-    ) -> Tuple[Tensor, Dim]:
+        out_dim: Dim,
+    ) -> Tensor:
         """
         Merges a list of axes into a single one. (Flatten the dims.)
         E.g. input is (batch, width, height, dim) and dims=(width,height), then we get (batch, width*height, dim).
@@ -251,18 +251,14 @@ class ReturnnLayersBackend(Backend[Layer]):
         :param source:
         :param dims:
         :param out_dim:
-        :return: tensor, out_dim
+        :return: tensor
         """
         if not isinstance(source, Tensor):
             raise TypeError(f"merge_dims: unexpected type for source {source!r}, need tensor")
-        if out_dim is None:
-            out_dim = dims[0]
-            for d in dims[1:]:
-                out_dim = out_dim * d
         layer = rfl.make_layer(
             {"class": "merge_dims", "from": source, "axes": dims, "out_dim": out_dim}, name="merge_dims"
         )
-        return layer, out_dim
+        return layer
 
     @staticmethod
     def split_dims(

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/returnn/torch/frontend/_backend.py

@@ -262,8 +262,8 @@ class TorchBackend(Backend[torch.Tensor]):
         source: Tensor,
         *,
         dims: Sequence[Dim],
-        out_dim: Optional[Dim] = None,
-    ) -> Tuple[Tensor, Dim]:
+        out_dim: Dim,
+    ) -> Tensor:
         """
         Merges a list of axes into a single one. (Flatten the dims.)
         E.g. input is (batch, width, height, dim) and dims=(width,height), then we get (batch, width*height, dim).
@@ -272,18 +272,12 @@ class TorchBackend(Backend[torch.Tensor]):
         :param source:
         :param dims:
         :param out_dim:
-        :return: tensor, out_dim
+        :return: tensor
         """
-        assert dims
-        if len(dims) == 1:
-            return source, dims[0]
+        assert len(dims) >= 2
         first_axis = min(source.dims.index(d) for d in dims)
         pre_dims = source.dims[:first_axis]
         post_dims = [d for d in source.dims if d not in dims and d not in pre_dims]
-        if out_dim is None:
-            out_dim = dims[0]
-            for d in dims[1:]:
-                out_dim = out_dim * d
         source = source.copy_transpose(tuple(pre_dims) + tuple(dims) + tuple(post_dims), allow_int=False)
         out = Tensor(
             "merge_dims",
@@ -295,7 +289,7 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.reshape(source.raw_tensor, out_shape)
         if source.feature_dim and source.feature_dim in dims:
             out.feature_dim = out_dim
-        return out, out_dim
+        return out
 
     @staticmethod
     def split_dims(

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/returnn/torch/util/debug_inf_nan.py

@@ -39,6 +39,7 @@ from io import TextIOBase
 import traceback
 from types import FrameType
 import torch
+import tree
 
 # noinspection PyProtectedMember
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -96,6 +97,7 @@ def debug_inf_nan(
 
 # For efficiency, and to be less spammy
 _TraceFuncNameBlacklist = {
+    "aten::empty.memory_format",
     "aten::zeros_like",
     "aten::ones_like",
     "aten::full",
@@ -113,6 +115,7 @@ _TraceFuncNameBlacklist = {
     "aten::split_with_sizes",
     "aten::slice.Tensor",
     "aten::select.int",
+    "aten::max_pool2d_with_indices",
 }
 
 
@@ -140,19 +143,20 @@ class _TraceOps(TorchDispatchMode):
         if self.report_every_op_call:
             print(f"--- op {func.name()}", file=self.file)
         out = func(*args, **kwargs)
-        if isinstance(out, torch.Tensor):
-            with no_python_dispatcher():
-                got_nan_inf_t = torch.stack([torch.isnan(out).any(), torch.isinf(out).any()]).cpu()
-                got_nan = got_nan_inf_t[0].item()
-                got_inf = got_nan_inf_t[1].item()
-                if got_nan or got_inf:
-                    s = "/".join([s_ for s_, b in [("nan", got_nan), ("inf", got_inf)] if b])
-                    print(f"--> {s} in {func}: {out}", file=self.file)
-                    traceback.print_list(
-                        _extract_stack_up_to(skip_top_num_frames=1, root_frame=self.root_frame), file=self.file
-                    )
-                    if self.stop_reporting_after_first_inf_nan:
-                        self.enabled = False
+        for out_ in tree.flatten(out):
+            if isinstance(out_, torch.Tensor):
+                with no_python_dispatcher():
+                    got_nan_inf_t = torch.stack([torch.isnan(out_).any(), torch.isinf(out_).any()]).cpu()
+                    got_nan = got_nan_inf_t[0].item()
+                    got_inf = got_nan_inf_t[1].item()
+                    if got_nan or got_inf:
+                        s = "/".join([s_ for s_, b in [("nan", got_nan), ("inf", got_inf)] if b])
+                        print(f"--> {s} in {func}: {out_}", file=self.file)
+                        traceback.print_list(
+                            _extract_stack_up_to(skip_top_num_frames=1, root_frame=self.root_frame), file=self.file
+                        )
+                        if self.stop_reporting_after_first_inf_nan:
+                            self.enabled = False
         return out
 
 

{returnn-1.20250225.201207 → returnn-1.20250226.115259/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250225.201207
+Version: 1.20250226.115259
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20250225.201207 → returnn-1.20250226.115259}/tests/rf_utils.py

@@ -17,6 +17,7 @@ from returnn.tensor import Tensor, Dim, TensorDict
 from returnn.tensor.utils import tensor_dict_fill_random_numpy_
 import returnn.torch.frontend as rft
 from returnn.torch.data.tensor_utils import tensor_dict_numpy_to_torch_, tensor_dict_torch_to_numpy_
+from returnn.torch.util.debug_inf_nan import debug_inf_nan as torch_debug_inf_nan
 
 # noinspection PyProtectedMember
 from returnn.frontend._random_journal import RandomJournal
@@ -72,8 +73,17 @@ def run_model(
         out_pt_raw = out_pt.as_raw_tensor_dict(include_const_sizes=True)
 
     if not allow_inf_nan_in_output:
+        non_finite_outputs = {}
         for k, v in out_pt.data.items():
-            assert numpy.isfinite(v.raw_tensor).all(), f"output {k!r} has non-finite values: {v.raw_tensor}"
+            if not numpy.isfinite(v.raw_tensor).all():
+                non_finite_outputs[k] = v
+                print(f"ERROR: output {k!r} has non-finite values:\n{v.raw_tensor}")
+        if non_finite_outputs:
+            torch_debug_inf_nan(
+                lambda: (_run_model_torch(extern_data, get_model, forward_step), None)[-1],
+                stop_reporting_after_first_inf_nan=False,
+            )
+            raise Exception(f"Non-finite values in output: {non_finite_outputs}. See log above.")
 
     if not test_tensorflow:
         return out_pt

returnn-1.20250225.201207/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250225.201207'
-long_version = '1.20250225.201207+git.c7cfe6c'