returnn 1.20240808.234227__tar.gz → 1.20240824.1611__tar.gz

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

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

returnn-1.20240824.1611/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240824.001611'
+long_version = '1.20240824.001611+git.a15f05e'

{returnn-1.20240808.234227 → returnn-1.20240824.1611}/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.20240824.1611}/returnn/frontend/array_.py

@@ -520,9 +520,12 @@ 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.
@@ -553,6 +556,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.20240824.1611}/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.20240824.1611}/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.20240824.1611}/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.20240824.1611}/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.20240824.1611}/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.20240824.1611}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240808.234227
+Version: 1.20240824.1611
 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.20240824.1611}/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'