returnn 1.20240808.234227__tar.gz → 1.20240826.132634__tar.gz

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/PKG-INFO

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

returnn-1.20240826.132634/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240826.132634'
+long_version = '1.20240826.132634+git.8b552f5'

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/frontend/_backend.py

@@ -1014,6 +1014,15 @@ class Backend(Generic[T]):
             return start + weight * (end - start)
         return rf.combine_bc(start, "+", rf.combine_bc(weight, "*", rf.combine_bc(end, "-", start)))
 
+    @staticmethod
+    def cumsum(source: Tensor, *, spatial_dim: Dim) -> Tensor:
+        """
+        :param source:
+        :param spatial_dim:
+        :return: cumsum over spatial dim
+        """
+        raise NotImplementedError
+
     @staticmethod
     def matmul(a: Tensor[T], b: Tensor[T], *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
         """

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/frontend/array_.py

@@ -520,19 +520,53 @@ def masked_select(
     tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None
 ) -> Tuple[Tensor, Dim]:
     """
+    This will pack the tensor based on the mask.
     In TF, this is ``boolean_mask``.
     The inverse of this is :func:`masked_scatter`.
 
+    Related: :func:`pack_padded`, which uses :func:`sequence_mask` as the mask.
+
     :param tensor:
     :param mask:
-    :param dims: the order of the dims defines the format. those dims should be exactly the dims of the mask.
-    :param out_dim:
-    :return: tensor where all dims in mask/dims are removed and replaced by a new dim.
-        the new dim is also returned.
-        if mask==True for all elements, the returned tensor would be simply the flattened input tensor.
-    """
-    # noinspection PyProtectedMember
-    return tensor._raw_backend.masked_select(tensor, mask=mask, dims=dims, out_dim=out_dim)
+    :param dims: The same or a subset of the dims of ``mask``.
+        The order of the dims defines the packed format (in ``out_dim``).
+        If only a single dim is given, only this dim is packed, while all other dims are kept as is.
+    :param out_dim: packed ``dims`` into this dim. Automatically created if not given, and then returned.
+    :return: Tensor where all ``dims`` are removed and replaced by a new dim (``out_dim``).
+        The new dim (``out_dim``) is also returned.
+        If mask==True for all elements, the returned tensor would be simply the flattened input tensor.
+    """
+    mask_dims_set = set(mask.dims)
+    dims_set = set(dims)
+    if not dims_set.issubset(mask_dims_set):
+        raise ValueError(f"masked_select: dims {dims} not subset of mask dims {mask.dims}")
+    if not dims_set:
+        raise ValueError(f"masked_select: dims {dims} empty")
+    if dims_set == mask_dims_set:
+        # noinspection PyProtectedMember
+        return tensor._raw_backend.masked_select(tensor, mask=mask, dims=dims, out_dim=out_dim)
+    # Separate implementation for the case where we have a subset of the mask dims, specifically one single dim.
+    # See https://github.com/rwth-i6/returnn/issues/1605 for discussion.
+    if len(dims) != 1:
+        # Please check https://github.com/rwth-i6/returnn/issues/1605 when you need this.
+        raise NotImplementedError(f"masked_select: dims {dims} len {len(dims)} > 1 not supported")
+    (in_dim,) = dims
+    in_dim: Dim
+    mask = mask.copy_masked(mask_value=False, dims=dims)
+    idxs = rf.cumsum(rf.cast(mask, "int32"), spatial_dim=in_dim)  # [T,B] -> idx in T' + 1
+    new_size = rf.gather(idxs, indices=in_dim.get_size_tensor() - 1, axis=in_dim)  # [B]
+    if out_dim is None:
+        out_dim = Dim(new_size, name="masked_select")
+    elif out_dim.dyn_size_ext is None:
+        out_dim.dyn_size_ext = new_size
+    elif out_dim.dyn_size_ext is not None and out_dim.dyn_size_ext.raw_tensor is None:
+        out_dim.dyn_size_ext.raw_tensor = new_size.raw_tensor
+    new_time = rf.reduce_max(new_size, axis=new_size.dims)  # T'
+    idxs = rf.where(mask, idxs - 1, new_time)  # new_time is the padding idx
+    ext_out_dim = out_dim + 1  # one more for the padded data
+    ext_res = rf.scatter(tensor, indices=idxs, indices_dim=in_dim, out_dim=ext_out_dim)
+    res, _ = rf.slice(ext_res, axis=ext_out_dim, size=out_dim)
+    return res, out_dim
 
 
 def masked_scatter(source: Tensor, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim) -> Tensor:
@@ -553,6 +587,7 @@ def sequence_mask(dims: Union[Dim, Sequence[Dim]], *, device: Optional[str] = No
     """
     :param dims:
     :param device:
+    :return: mask based on the sequence lengths
     """
     if isinstance(dims, Dim):
         dims = [dims]

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/frontend/math_.py

@@ -69,6 +69,7 @@ __all__ = [
     "log_softmax",
     "gating",
     "lerp",
+    "cumsum",
 ]
 
 
@@ -612,3 +613,11 @@ def lerp(
     """
     # noinspection PyProtectedMember
     return start._raw_backend.lerp(start, end, weight, allow_broadcast_all_sources=allow_broadcast_all_sources)
+
+
+def cumsum(source: Tensor, *, spatial_dim: Dim) -> Tensor:
+    """
+    Applies cumsum.
+    """
+    # noinspection PyProtectedMember
+    return source._raw_backend.cumsum(source, spatial_dim=spatial_dim)

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/tf/frontend_layers/_backend.py

@@ -730,6 +730,11 @@ class ReturnnLayersBackend(Backend[Layer]):
             name="clip_by_value",
         )
 
+    @staticmethod
+    def cumsum(source: Tensor, *, spatial_dim: Dim) -> Tensor:
+        """cumsum"""
+        return rfl.make_layer({"class": "cumsum", "from": source, "axis": spatial_dim}, name="cumsum")
+
     @staticmethod
     def matmul(a: Tensor, b: Tensor, *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor:
         """matmul"""

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/torch/frontend/_backend.py

@@ -1192,6 +1192,14 @@ class TorchBackend(Backend[torch.Tensor]):
         )
         return out
 
+    @staticmethod
+    def cumsum(source: Tensor, *, spatial_dim: Dim) -> Tensor:
+        """cumsum"""
+        axis = source.get_axis_from_description(spatial_dim)
+        out = source.copy_template("cumsum")
+        out.raw_tensor = torch.cumsum(source.raw_tensor, dim=axis, dtype=source.raw_tensor.dtype)
+        return out
+
     @staticmethod
     def matmul(a: _TT, b: _TT, *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> _TT:
         """

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/torch/updater.py

@@ -79,7 +79,7 @@ def _get_class_init_kwargs(optim_class):
     return optim_class_init_kwargs
 
 
-class Updater(object):
+class Updater:
     """
     Wraps a torch.optim.Optimizer, and extends it by some further functionality.
     """

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn/util/debug.py

@@ -3,8 +3,8 @@ Some generic debugging utilities.
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Union, Any, Sequence, Callable, Tuple, List, Dict
-from types import FunctionType
+from typing import TYPE_CHECKING, Optional, Union, Any, Collection, Sequence, Callable, Tuple, List, Dict, TextIO
+from types import FunctionType, CodeType, FrameType
 import os
 import sys
 import signal
@@ -54,13 +54,19 @@ def auto_exclude_all_new_threads(func):
     return wrapped
 
 
-def dump_all_thread_tracebacks(exclude_thread_ids=None, exclude_self=False):
+def dump_all_thread_tracebacks(
+    *, exclude_thread_ids: Optional[Collection[int]] = None, exclude_self: bool = False, file: Optional[TextIO] = None
+):
     """
-    :param set[int] exclude_thread_ids:
-    :param bool exclude_self:
+    :param exclude_thread_ids:
+    :param exclude_self:
+    :param file:
     """
     if exclude_thread_ids is None:
         exclude_thread_ids = set()
+    if file is None:
+        file = sys.stdout
+
     from returnn.util.better_exchook import print_tb
     import threading
 
@@ -68,7 +74,7 @@ def dump_all_thread_tracebacks(exclude_thread_ids=None, exclude_self=False):
         exclude_thread_ids = set(list(exclude_thread_ids) + [threading.current_thread().ident])
 
     if hasattr(sys, "_current_frames"):
-        print("")
+        print("", file=file)
         threads = {t.ident: t for t in threading.enumerate()}
         # noinspection PyProtectedMember
         for tid, stack in sorted(sys._current_frames().items()):
@@ -89,17 +95,17 @@ def dump_all_thread_tracebacks(exclude_thread_ids=None, exclude_self=False):
                 tags += [str(thread_)]
             else:
                 tags += ["unknown with id %i" % tid]
-            print("Thread %s:" % ", ".join(tags))
+            print("Thread %s:" % ", ".join(tags), file=file)
             if tid in global_exclude_thread_ids:
-                print("(Auto-ignored traceback.)")
+                print("(Auto-ignored traceback.)", file=file)
             elif tid in exclude_thread_ids:
-                print("(Excluded thread.)")
+                print("(Excluded thread.)", file=file)
             else:
-                print_tb(stack, file=sys.stdout)
-            print("")
-        print("That were all threads.")
+                print_tb(stack, file=file)
+            print("", file=file)
+        print("That were all threads.", file=file)
     else:
-        print("Does not have sys._current_frames, cannot get thread tracebacks.")
+        print("Does not have sys._current_frames, cannot get thread tracebacks.", file=file)
 
 
 def setup_warn_with_traceback():
@@ -132,6 +138,7 @@ def init_better_exchook():
     but adds some special handling for the main thread.
     """
     from returnn.util.better_exchook import better_exchook
+    from returnn.log import log
 
     def excepthook(exc_type, exc_obj, exc_tb):
         """
@@ -139,6 +146,8 @@ def init_better_exchook():
         :param exc_obj:
         :param exc_tb:
         """
+        file = log.v1 or sys.stdout
+
         # noinspection PyBroadException
         try:
             # noinspection PyUnresolvedReferences,PyProtectedMember
@@ -153,7 +162,10 @@ def init_better_exchook():
                     return
                 # An unhandled exception in the main thread. This means that we are going to quit now.
                 sys.exited = True
-        print("Unhandled exception %s in thread %s, proc %i." % (exc_type, threading.current_thread(), os.getpid()))
+        print(
+            "Unhandled exception %s in thread %s, proc %i." % (exc_type, threading.current_thread(), os.getpid()),
+            file=file,
+        )
         if exc_type is KeyboardInterrupt:
             return
 
@@ -164,9 +176,9 @@ def init_better_exchook():
                 # We are the main thread and we got an exit-exception. This is likely fatal.
                 # This usually means an exit. (We ignore non-daemon threads and procs here.)
                 # Print the stack of all other threads.
-                dump_all_thread_tracebacks(exclude_thread_ids={main_thread_id})
+                dump_all_thread_tracebacks(exclude_thread_ids={main_thread_id}, file=file)
 
-        better_exchook(exc_type, exc_obj, exc_tb, file=sys.stdout)
+        better_exchook(exc_type, exc_obj, exc_tb, file=file)
 
     sys.excepthook = excepthook
 
@@ -574,6 +586,15 @@ class PyTracer:
         check_py_traces_rf_to_pt_equal(trace_my_impl.captured_locals, trace_ref_impl.captured_locals, [...])
 
     See also :func:`check_py_traces_rf_to_pt_equal` to compare the traces.
+
+    This class uses the Python :func:`sys.settrace` mechanism to trace the locals.
+    It accesses ``frame.f_locals`` to get the local variables.
+    Note that this behavior is slightly buggy in versions of CPython <3.13,
+    see for example:
+    https://github.com/python/cpython/issues/113939
+    https://github.com/python/cpython/issues/74929
+    And thus the behavior might be different depending on the Python version.
+    In Python >=3.13, you likely get a few more locals than before.
     """
 
     def __init__(
@@ -583,13 +604,14 @@ class PyTracer:
         :param funcs_to_trace_list: list of functions to trace the locals. only those functions will be traced.
         :param capture_type: only capture variables of this type, e.g. torch.Tensor.
         """
-        self.funcs_to_trace_list = []
-        for func in funcs_to_trace_list:
-            # E.g. when the function was decorated, see functools.wraps.
+
+        def _get_func_code(func: FunctionType) -> CodeType:
             while getattr(func, "__wrapped__", None) is not None:
                 func = func.__wrapped__
-            self.funcs_to_trace_list.append(func)
-        self._code_obj_to_func = {func.__code__: func for func in self.funcs_to_trace_list}
+            return func.__code__
+
+        self.funcs_to_trace_list = funcs_to_trace_list
+        self._code_obj_to_func = {_get_func_code(func): func for func in self.funcs_to_trace_list}
         self.capture_type = capture_type
 
         self._prev_trace_func = None
@@ -611,7 +633,7 @@ class PyTracer:
         sys.settrace(self._prev_trace_func)
         self._prev_trace_func = None
 
-    def __call__(self, frame, event, arg) -> Optional[PyTracer]:
+    def __call__(self, frame: FrameType, event, arg) -> Optional[PyTracer]:
         """
         Trace func to get intermediate outputs.
         """
@@ -623,16 +645,19 @@ class PyTracer:
             if event == "call":
                 self.captured_locals.setdefault(func, []).append({})
             else:
+                # Note that accessing frame.f_locals is not always totally accurate.
+                # See the corresponding comment in our class docstring.
                 for k, v in frame.f_locals.items():
                     if not isinstance(v, self.capture_type):
                         continue
                     prev = self.captured_locals[func][-1].get(k, None)
                     if prev is None or prev[-1] is not v:
+                        var_ls = self.captured_locals[func][-1].setdefault(k, [])
                         print(
                             f"{func.__qualname__}[{len(self.captured_locals[func]) - 1}]"
-                            f" {type(v).__qualname__} var changed: {k} = {v}"
+                            f" ({type(v).__qualname__}) {k}[{len(var_ls)}] = {v}"
                         )
-                        self.captured_locals[func][-1].setdefault(k, []).append(v)
+                        var_ls.append(v)
             return self
         return prev_trace_func_res
 
@@ -675,8 +700,6 @@ def check_py_traces_rf_to_pt_equal(
     dummy_forward_compat_kwargs = {f"_random_forward_compat_arg_{random.randint(0, 1_000_000)}": "please"}
 
     def _get_entry(trace, func, i, name, j):
-        while getattr(func, "__wrapped__", None) is not None:
-            func = func.__wrapped__
         return trace[func][i][name][j]
 
     def _resolve_dim(dim: Union[Dim, str]) -> Dim:

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/tests/test_rf_array.py

@@ -33,6 +33,76 @@ def test_pack_padded():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_masked_select():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "mask": Tensor("mask", [batch_dim, time_dim], dtype="bool"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x = extern_data["data"]
+        mask = extern_data["mask"]
+        out, pack_dim = rf.masked_select(x, mask=mask, dims=[batch_dim, time_dim])
+        out.mark_as_default_output(shape=(pack_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_masked_select_single_dim():
+    # https://github.com/rwth-i6/returnn/issues/1605
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "mask": Tensor("mask", [batch_dim, time_dim], dtype="bool"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x = extern_data["data"]
+        mask = extern_data["mask"]
+        out, pack_dim = rf.masked_select(x, mask=mask, dims=[time_dim])
+        out.mark_as_default_output(shape=(batch_dim, pack_dim, in_dim))
+        pack_dim.dyn_size_ext.mark_as_output("out_size", shape=[batch_dim])
+        x.mark_as_output("input", shape=[batch_dim, time_dim, in_dim])
+        mask.mark_as_output("mask", shape=[batch_dim, time_dim])
+        time_dim.dyn_size_ext.mark_as_output("in_size", shape=[batch_dim])
+
+    res = run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step, test_tensorflow=False)
+    for k, v in res.data.items():
+        print(f"{k} {v} = {v.raw_tensor}")
+    num_batches = res["input"].raw_tensor.shape[0]
+    have_non_zero = False
+    have_less = False
+    for b in range(num_batches):
+        in_size = res["in_size"].raw_tensor[b]
+        out_size = res["out_size"].raw_tensor[b]
+        print(f"batch {b}, in_size {in_size}, out_size {out_size}")
+        assert 0 <= out_size <= in_size
+        if out_size > 0:
+            have_non_zero = True
+        if out_size < in_size:
+            have_less = True
+        c = 0
+        for t in range(in_size):
+            if res["mask"].raw_tensor[b, t]:
+                c += 1
+        assert c == out_size
+        t_ = 0
+        for t in range(in_size):
+            if res["mask"].raw_tensor[b, t]:
+                assert (res["input"].raw_tensor[b, t] == res["output"].raw_tensor[b, t_]).all()
+                t_ += 1
+    assert have_non_zero and have_less  # just that the test case covered all cases
+
+
 def test_reshape():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

{returnn-1.20240808.234227 → returnn-1.20240826.132634}/tests/test_rf_math.py

@@ -109,3 +109,19 @@ def test_log_sigmoid():
         out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
 
     run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_cumsum():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim], dtype="int32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: rf.Module, extern_data: TensorDict):
+        out = rf.cumsum(extern_data["data"], spatial_dim=time_dim)
+        out.mark_as_default_output(shape=(batch_dim, time_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)

returnn-1.20240808.234227/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240808.234227'
-long_version = '1.20240808.234227+git.d773e45'