returnn 1.20250429.161207__py3-none-any.whl → 1.20250508.93313__py3-none-any.whl

returnn/PKG-INFO

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

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20250429.161207'
-long_version = '1.20250429.161207+git.d4d26bf'
+version = '1.20250508.093313'
+long_version = '1.20250508.093313+git.4f05ac7'

returnn/frontend/array_.py

@@ -39,6 +39,7 @@ __all__ = [
     "pad_packed",
     "gather",
     "scatter",
+    "scatter_mean",
     "scatter_argmax",
     "scatter_logsumexp",
     "scatter_logmeanexp",
@@ -807,8 +808,8 @@ def scatter(
     :param source: [batch_dims..., indices_dim(s)..., feature_dims...]
     :param indices: [batch_dims..., indices_dim(s)...] -> out_dim
     :param indices_dim:
-    :param mode: "sum", "max", "min", "logsumexp", "logmeanexp", "argmax".
-        (Note: If you ever need mean, argmin, etc, please open an issue/PR.)
+    :param mode: "sum", "max", "min", "mean", "logsumexp", "logmeanexp", "argmax".
+        (Note: If you ever need another mode, please open an issue/PR.)
     :param fill_value:
     :param out_dim: The indices target dim.
         If not given, will be automatically determined as the sparse_dim from indices.
@@ -817,6 +818,8 @@ def scatter(
     :param use_mask:
     :return: [batch_dims..., out_dim(s)..., feature_dims...]
     """
+    if mode == "mean":
+        return scatter_mean(source, indices=indices, indices_dim=indices_dim, fill_value=fill_value, out_dim=out_dim)
     if mode == "logsumexp":
         return scatter_logsumexp(
             source, indices=indices, indices_dim=indices_dim, fill_value=fill_value, out_dim=out_dim
@@ -863,6 +866,35 @@ def scatter(
     return out
 
 
+def scatter_mean(
+    source: Tensor,
+    *,
+    indices: Tensor,
+    indices_dim: Union[Dim, Sequence[Dim]],
+    fill_value: Optional[Union[int, float]] = None,
+    out_dim: Optional[Union[Dim, Sequence[Dim]]] = None,
+) -> Tensor:
+    """
+    Scatters into new zero-tensor.
+    If entries in indices are duplicated, the corresponding values in source will be mean'ed together.
+    This is like :func:`scatter` with ``mode="mean"``.
+
+    :param source: [batch_dims..., indices_dim(s)..., feature_dims...]
+    :param indices: [batch_dims..., indices_dim(s)...] -> out_dim
+    :param indices_dim:
+    :param fill_value:
+    :param out_dim: The indices target dim.
+        If not given, will be automatically determined as the sparse_dim from indices.
+        If multiple out dims, use indices into the merged out dims,
+        and then we use :func:`rf.split_dims` afterwards.
+    :return: [batch_dims..., out_dim(s)..., feature_dims...]
+    """
+    ones = rf.ones(dims=indices.dims, dtype=source.dtype, device=source.device)
+    counts = rf.scatter(ones, indices=indices, indices_dim=indices_dim, fill_value=1, out_dim=out_dim)
+    y = scatter(source, indices=indices, indices_dim=indices_dim, fill_value=fill_value, out_dim=out_dim)
+    return y / counts
+
+
 def scatter_argmax(
     source: Tensor,
     *,

returnn/util/better_exchook.py

@@ -58,6 +58,7 @@ import threading
 import keyword
 import inspect
 import contextlib
+from weakref import WeakKeyDictionary
 
 try:
     import typing
@@ -1564,6 +1565,9 @@ def get_func_str_from_code_object(co, frame=None):
     return co.co_name
 
 
+_func_from_code_object_cache = WeakKeyDictionary()  # code object -> function
+
+
 def get_func_from_code_object(co, frame=None):
     """
     :param types.CodeType co:
@@ -1580,6 +1584,11 @@ def get_func_from_code_object(co, frame=None):
     import types
 
     assert isinstance(co, (types.CodeType, DummyFrame))
+    co_is_code_object = isinstance(co, types.CodeType)
+    if co_is_code_object:
+        candidate = _func_from_code_object_cache.get(co)
+        if candidate:
+            return candidate
     _attr_name = "__code__" if PY3 else "func_code"
     if frame and frame.f_code.co_nlocals > 0:
         func_name = frame.f_code.co_name
@@ -1587,18 +1596,23 @@ def get_func_from_code_object(co, frame=None):
         if frame_self is not None:
             candidate = getattr(frame_self.__class__, func_name, None)
             if candidate and (getattr(candidate, _attr_name, None) is co or isinstance(co, DummyFrame)):
+                if co_is_code_object:
+                    _func_from_code_object_cache[co] = candidate
                 return candidate
     try:
         candidate = getattr(_get_loaded_module_from_filename(co.co_filename), co.co_name, None)
     except ImportError:  # some modules have lazy loaders, but those might fail here
         candidate = None
     if candidate and (getattr(candidate, _attr_name, None) is co or isinstance(co, DummyFrame)):
+        if co_is_code_object:
+            _func_from_code_object_cache[co] = candidate
         return candidate
     if isinstance(co, DummyFrame):
         return None
     candidates = gc.get_referrers(co)
     candidates = [f for f in candidates if getattr(f, _attr_name, None) is co]
     if candidates:
+        _func_from_code_object_cache[co] = candidates[0]
         return candidates[0]
     return None
 

{returnn-1.20250429.161207.dist-info → returnn-1.20250508.93313.dist-info}/METADATA

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

{returnn-1.20250429.161207.dist-info → returnn-1.20250508.93313.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=HyJ82YwGT_Vw0szcqg-a4TqmDp-OQ2YsRXQ8rM5gTg8,5215
+returnn/PKG-INFO,sha256=Rsc8t3mGIL2V2iJflkOuNthzptMNFK9cMC-KiImaBDg,5214
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=qBFbuB1yN3adgVM5pXt2-Yq9vorjRNchNPL8kDKx44M,31752
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=DR3lkkJjhQ4JxpXrXSh3Yr0jKAWL6KbWNKJUE9DMJPs,77
+returnn/_setup_info_generated.py,sha256=cYPqIVcVewY03CWaIMT3pOmIFS88wUWDSj14rjsgy3Q,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -80,7 +80,7 @@ returnn/frontend/_cache.py,sha256=JAhi7L-raQ3A-NC3JUYDtdRTwT3BGJJGGZxrZ8MfEWQ,84
 returnn/frontend/_numpy_backend.py,sha256=fZjks7p3dgxVZ6tSDazTTgBxNjJqXjfqgw_7mA7rDEE,9066
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
 returnn/frontend/_utils.py,sha256=4A3MSRM0i86J77550uR_AjcBEPu6nymLUZ9Xd1V3Fkc,12073
-returnn/frontend/array_.py,sha256=eYwH-NVAoHpVrFdJv08lCqh3jvfoZV_ZBEoWHjsBz0o,50090
+returnn/frontend/array_.py,sha256=o_NSq87pB5I2XvFUjk40Dobqx6tTfEY1wzgmaelujgM,51511
 returnn/frontend/attention.py,sha256=GKt-Xqnz8sIyXVrE0i4VCS7J2Wu7dmoH_BA0Cu8CrXQ,45769
 returnn/frontend/backend.py,sha256=iQ9w4xl8Ea7bgpb0VUaCKq50rV5Bl2E5J8Rhd-oqD_c,883
 returnn/frontend/build_from_dict.py,sha256=rfWa2rjjhIR_kIQED_nMrygrQBunS6unegzWTLVbC98,3017
@@ -234,7 +234,7 @@ returnn/torch/util/module.py,sha256=MXHIrF9Isu575DDJIa81212ULKwdqu1oOLxDVZecVSk,
 returnn/torch/util/scaled_gradient.py,sha256=3585VuNypBty-pW6r3BKK047H3MqZQSdMjXeYAb4cmU,3192
 returnn/util/__init__.py,sha256=UIG1qw4idqhW71BV60ha7h9PktxvEVcBIu0lYRossK8,336
 returnn/util/basic.py,sha256=jwtaaZyOV7fUjKDRXVHDy-K5kwR1mPrkAZrzc5STOvE,142554
-returnn/util/better_exchook.py,sha256=TAtb_ZyM-357UnOg_HMoBZUSxzt0WPgumlvprmlCprA,63921
+returnn/util/better_exchook.py,sha256=98XnUZIWpYN7NfklSGt_5hYNplADVFQnh857esKxjdI,64475
 returnn/util/bpe.py,sha256=LWFhICZsEOnMwNws0lybPNzKRX6rSr8yKCvP65vjl9Y,19656
 returnn/util/debug.py,sha256=wuRzdg9zB84WWCGyTjmRR_zYypu8gXxlc0nZ6si9OC8,28224
 returnn/util/debug_helpers.py,sha256=0EINLK4uLtoSt5_kHs1M2NIFpMd0S7i4c4rx90U4fJk,2914
@@ -253,8 +253,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20250429.161207.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250429.161207.dist-info/METADATA,sha256=HyJ82YwGT_Vw0szcqg-a4TqmDp-OQ2YsRXQ8rM5gTg8,5215
-returnn-1.20250429.161207.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20250429.161207.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250429.161207.dist-info/RECORD,,
+returnn-1.20250508.93313.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250508.93313.dist-info/METADATA,sha256=Rsc8t3mGIL2V2iJflkOuNthzptMNFK9cMC-KiImaBDg,5214
+returnn-1.20250508.93313.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20250508.93313.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250508.93313.dist-info/RECORD,,