returnn 1.20250701.140328__py3-none-any.whl → 1.20250703.183400__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250701.140328
+Version: 1.20250703.183400
 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.20250701.140328'
-long_version = '1.20250701.140328+git.6646903'
+version = '1.20250703.183400'
+long_version = '1.20250703.183400+git.a5c35a2'

returnn/frontend/_backend.py

@@ -220,7 +220,8 @@ class Backend(Generic[T]):
         """
         :param a:
         :param kind: "add", "sub", "mul", "truediv", "floordiv", "mod", "pow",
-            "maximum", "minimum", "logical_and", "logical_or", "squared_difference"
+            "maximum", "minimum", "logical_and", "logical_or", "squared_difference",
+            "logaddexp"
         :param b:
         :return: a `kind` b
         """

returnn/frontend/_native/module.cpp

@@ -316,6 +316,7 @@ bool PyModuleState::_cachedOpInitTorch() {
     AddOp(TOp_Maximum, "clamp_min");
     AddOp(TOp_Minimum, "clamp_max");
     AddOpAlt(TOp_SquaredDifference, "squared_difference");
+    AddOp(TOp_LogAddExp, "logaddexp");
     AddOp(TOp_And, "logical_and");
     AddOp(TOp_Or, "logical_or");
     AddOp(TOp_Neg, "neg");
@@ -356,6 +357,7 @@ const char* rawOpName(RawOp op) {
         names[TOp_Maximum] = "maximum";
         names[TOp_Minimum] = "minimum";
         names[TOp_SquaredDifference] = "squared_difference";
+        names[TOp_LogAddExp] = "logaddexp";
         names[TOp_And] = "logical_and";
         names[TOp_Or] = "logical_or";
         // The names for the unary funcs matter:

returnn/frontend/_native/module.hpp

@@ -34,6 +34,7 @@ enum RawOp {
     TOp_Maximum,
     TOp_Minimum,
     TOp_SquaredDifference,
+    TOp_LogAddExp,
 
     TOp_And,
     TOp_Or,

returnn/frontend/_native/tensor_ops.cpp

@@ -1468,6 +1468,7 @@ static PyObject* _pyTensorCompareOrCombine(PyObject *self, PyObject *args, PyObj
         kindToCombineFunc["logical_and"] = TOp_And;
         kindToCombineFunc["logical_or"] = TOp_Or;
         kindToCombineFunc["squared_difference"] = TOp_SquaredDifference;
+        kindToCombineFunc["logaddexp"] = TOp_LogAddExp;
     }
 
     auto it = isCompare ? kindToCompareFunc.find(kind) : kindToCombineFunc.find(kind);

returnn/frontend/math_.py

@@ -37,6 +37,7 @@ __all__ = [
     "logical_not",
     "opt_logical_or",
     "opt_logical_and",
+    "log_add_exp",
     "is_finite",
     "is_infinite",
     "is_neg_infinite",
@@ -173,7 +174,8 @@ def combine(
     """
     :param a:
     :param kind: "add"|"+", "sub"|"-", "mul"|"*", "truediv"|"/", "floordiv"|"//", "mod"|"%", "pow"|"**",
-        "max"|"maximum", "min"|"minimum", "logical_and", "logical_or", "squared_difference"
+        "max"|"maximum", "min"|"minimum", "logical_and", "logical_or", "squared_difference",
+        "logaddexp"
     :param b:
     :param allow_broadcast_all_sources: if True, it is allowed that neither a nor b has all dims of the result.
         Not needed when out_dims is specified explicitly.
@@ -364,6 +366,16 @@ def opt_logical_and(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Ten
     return combine(a, "logical_and", b)
 
 
+def log_add_exp(a: Tensor, b: Tensor) -> Tensor:
+    """
+    Computes log(exp(a) + exp(b)) in a numerically stable way.
+    This is useful for log probabilities, e.g. in beam search.
+
+    See also: func:`reduce_logsumexp`.
+    """
+    return combine(a, "logaddexp", b)
+
+
 def is_finite(a: Tensor) -> Tensor:
     """is finite"""
     # noinspection PyProtectedMember

returnn/frontend/reduce.py

@@ -99,6 +99,8 @@ def reduce_logsumexp(source: Tensor[T], *, axis: Union[Dim, Sequence[Dim]], use_
     """
     Reduce the tensor along the given axis
 
+    Also see :func:`log_add_exp`.
+
     :param source:
     :param axis:
     :param use_mask: if True (default), use the time mask (part of dim tag) to ignore padding frames
@@ -247,7 +249,7 @@ class RunningMean(rf.Module):
         """
 
         def _update_running_stats():
-            assert all(d in self.shape for d in x.dims)
+            assert all(d in x.dims for d in self.shape)
             x_ = rf.reduce_mean(x, axis=[d for d in x.dims if d not in self.shape])
             self.mean.assign_add(self.alpha * (x_ - self.mean))
 

returnn/tf/layers/basic.py

@@ -9607,6 +9607,7 @@ class CombineLayer(LayerBase):
           `maximum`, `minimum`,
           `logical_and`, `logical_or`,
           `squared_difference`,
+          `logaddexp`,
           or `eval`,
           or any function in the tf.math or tf namespace.
         :param list[LayerBase] sources:
@@ -9814,6 +9815,10 @@ class CombineLayer(LayerBase):
             assert kind == "eval" and eval_str
             return self._op_kind_eval(sources, eval_str=eval_str, eval_locals=eval_locals)
 
+        if hasattr(self, "_op_kind_%s" % kind):
+            func = getattr(self, "_op_kind_%s" % kind)
+            return func(sources)
+
         kind = {
             "+": "add",
             "-": "subtract",
@@ -9823,16 +9828,18 @@ class CombineLayer(LayerBase):
             "sub": "subtract",
             "mul": "multiply",
         }.get(kind, kind)
+
         if hasattr(tf, "math") and hasattr(tf.math, kind):
             tf_func = getattr(tf.math, kind)
         elif hasattr(tf, kind):
             tf_func = getattr(tf, kind)
+        elif hasattr(tf, "keras") and hasattr(tf.keras, "ops") and hasattr(tf.keras.ops, kind):
+            tf_func = getattr(tf.keras.ops, kind)
+        elif hasattr(tf, "experimental") and hasattr(tf.experimental, "numpy") and hasattr(tf.experimental.numpy, kind):
+            tf_func = getattr(tf.experimental.numpy, kind)
         else:
-            tf_func = None
-        if tf_func:
-            return self._op_dense_fn(sources, tf_func, self.output)
-
-        return getattr(self, "_op_kind_%s" % kind)(sources)
+            raise ValueError(f"{self}: unknown kind {kind!r}")
+        return self._op_dense_fn(sources, tf_func, self.output)
 
 
 class EvalLayer(CombineLayer):
@@ -10657,7 +10664,7 @@ class SearchSortedLayer(LayerBase):
         transposed_values_data = values_data.copy_transpose(perm=values_batch_axes + values_non_batch_axes)  # [B,F]
         x = transposed_sorted_data.placeholder  # [B,T]
         if transposed_sorted_data.dims[-1].need_masking():
-            from returnn.tf.util.basic import where_bc, sequence_mask
+            from returnn.tf.util.basic import where_bc
 
             seq_mask = transposed_sorted_data.get_sequence_mask_broadcast(axis=-1)
             x = where_bc(seq_mask, x, x.dtype.max)  # note: this is not correct if values contains x.dtype.max

returnn/torch/frontend/_backend.py

@@ -840,7 +840,8 @@ class TorchBackend(Backend[torch.Tensor]):
         """
         :param a:
         :param kind: "add", "sub", "mul", "truediv", "floordiv", "mod", "pow",
-            "maximum", "minimum", "logical_and", "logical_or", "squared_difference"
+            "maximum", "minimum", "logical_and", "logical_or", "squared_difference",
+            "logaddexp"
         :param b:
         :return: a `kind` b
         """

{returnn-1.20250701.140328.dist-info → returnn-1.20250703.183400.dist-info}/METADATA

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

{returnn-1.20250701.140328.dist-info → returnn-1.20250703.183400.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=RByOzAa8sFb4gdGsaWiRUEOlZq3X9qqn4MEF0ziNPNc,5215
+returnn/PKG-INFO,sha256=6b050edAnN_IixiTU7cKY9Y7GsV2u54Go4pMebtP1LA,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=lHyZcu_0yc9f7Vf_Kfdy9PmeU0T76XVXnpalHi5WKro,31740
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=FycWE5HycEgzbwEzgqhdpFdUBY_4Kg7Rlg43mZCjiBs,77
+returnn/_setup_info_generated.py,sha256=W3RqBPspEAo2psz8RlzIp3A0pBAmerlhlAfLeZuEUy8,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
@@ -75,7 +75,7 @@ returnn/extern/graph_editor/subgraph.py,sha256=q9o0zVBLDrTIidaXg5WG5daDW0mLbwv2J
 returnn/extern/graph_editor/transform.py,sha256=qMGSenpbAnGqdG6QP6iWjlm6_ccySYJaZKOoAj1dbOM,29348
 returnn/extern/graph_editor/util.py,sha256=HfRbyQPmQ6_n5-O-096n0KeJtllQXFtaurpeJS_URZ0,18706
 returnn/frontend/__init__.py,sha256=2aS7nbxXniIrBp2DODl0xN0f3IJ_dX4Bi9ZlR7W5_DE,1472
-returnn/frontend/_backend.py,sha256=w-Xtn-2-uu0AGL0VAjP5QFuh_tpqX2FiSyLkPvZ5bPY,50446
+returnn/frontend/_backend.py,sha256=pAnVAbZhIGKD-10tp0Mx7AO1GZNghYu7AVAPhiimN-k,50471
 returnn/frontend/_cache.py,sha256=JAhi7L-raQ3A-NC3JUYDtdRTwT3BGJJGGZxrZ8MfEWQ,8403
 returnn/frontend/_numpy_backend.py,sha256=fZjks7p3dgxVZ6tSDazTTgBxNjJqXjfqgw_7mA7rDEE,9066
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
@@ -101,7 +101,7 @@ returnn/frontend/label_smoothing.py,sha256=lxmaowNr61sCMzMewqHhu1r0CcklYfhLXlFnB
 returnn/frontend/linear.py,sha256=xRUjnkD3MTWDezSaYATBYJQ2fa1RhKMNrTuhC54hhVs,2252
 returnn/frontend/loop.py,sha256=t-z6ke1X03I2aPUEqLYmVZWyMzfW3IedFvKUGc-TCX8,16160
 returnn/frontend/loss.py,sha256=uSvou2MPd13JiLAg_OIQ3AyyLvD3RHjMEVgFEN0gKqU,7440
-returnn/frontend/math_.py,sha256=KlJxdIib8ENlid7cc4lcwHv5e21tzTjTEV8VgEDAijo,16984
+returnn/frontend/math_.py,sha256=A_RkZ5lH2uXMchfPIH3itraWtMNNCVckQHHpf7aIIZQ,17295
 returnn/frontend/matmul.py,sha256=xkueyxzSDz8MsYaWxPSjmV2Yy-tcaiOQDXbFt1IQM2A,1944
 returnn/frontend/module.py,sha256=219rh5mE0CD0-NdxXLsKyhv3BNtOI9jSyiI1Rb8MOyU,10700
 returnn/frontend/nested.py,sha256=P84u_cjoYdYRJ_0Cbt0vlKXxskmXTDfsnw_vFCCNKtU,15107
@@ -112,7 +112,7 @@ returnn/frontend/parametrize.py,sha256=VhgTEP7ehON950Q4bkCy8rvg9641moEKAXn0XzomK
 returnn/frontend/piecewise_linear.py,sha256=TdL6wzop8P1dcIZwkEbJFvSUZSI1cbhS3XKzlWQkEVI,1964
 returnn/frontend/rand.py,sha256=Levgf5VtOOBKDSgz0869Jf3VW4BWxYZuRXsa_fOxNI4,12969
 returnn/frontend/rec.py,sha256=6YSsSG7fdtfvvg24vmexSg8R2aVCcKHBdGLh-Mgn9Co,8037
-returnn/frontend/reduce.py,sha256=xvxN_h3LsMJdmT0IbW4nOf8qFhckuAniIhD9PalO6j0,10305
+returnn/frontend/reduce.py,sha256=gRSvBJZNHa757IqBxGw4hu5eiO3pjie_ptEwUXHLSCs,10340
 returnn/frontend/run_ctx.py,sha256=yyOMUCKTOe19C4z2Nfly4YCLBmQ9ihip6nGrkW-Y6qg,23789
 returnn/frontend/signal.py,sha256=hfDipDhO0n9nXhGy7txwYUNbvg28NqkFq9p0Jq46f9c,4411
 returnn/frontend/state.py,sha256=EePdrx6PtWL4mJ2XZmGlh5dl4nq6G9wZpqP4hdDEzfY,2935
@@ -122,10 +122,10 @@ returnn/frontend/types.py,sha256=r-QsxPQyFSr9WwCRzqTn_X5jQLbjthrtjHavY8XIDmk,109
 returnn/frontend/_native/__init__.py,sha256=fVjazAujt0rdICXZL-GgW1sjFeL1HB4NPuy2m5rmMsc,6480
 returnn/frontend/_native/backend.cpp,sha256=MeHczHypwj_ncntOxRqanK8SqGyV9Eq1X0cpMWb_WII,4768
 returnn/frontend/_native/backend.hpp,sha256=Wq80dcEzXfRNxGOXFnIgHllkiv1rDi3KpHK-xxJsSDI,791
-returnn/frontend/_native/module.cpp,sha256=lS1Oypo3n6oCu6cxKAmqpNjSvQN9aMZIOeMec96FWYU,15626
-returnn/frontend/_native/module.hpp,sha256=uf4HPSTrFP2brGR_x9G5N1ZlZ-ok5GakMbNo4LbqxUg,6670
+returnn/frontend/_native/module.cpp,sha256=9BCUoDTZDJ6hlXp4pUus1BlN7-oxcRy6tK9ctyCkwk0,15709
+returnn/frontend/_native/module.hpp,sha256=iv4jvQidLaE8uC-YbaYjiXONTL_Pq7WUQKQ5MdFpdIs,6689
 returnn/frontend/_native/py_utils.hpp,sha256=vcxKGmOyDRuwsmmSEjoaCJyKMy1BNYoGlso2pZu7VoE,3139
-returnn/frontend/_native/tensor_ops.cpp,sha256=bYtwwn_NeJfAEHWYPEJlkoLDKt9baZ3RA8av7gtz2qc,70246
+returnn/frontend/_native/tensor_ops.cpp,sha256=bA4Gf-q8cVENL441r1IYVd44EcUsV-eELyDzqmnCuw0,70302
 returnn/frontend/_native/tensor_ops.hpp,sha256=dDqvUejRNHjItnmOP5aHyAQbAmXmXoDVXSe3tveEU8A,3732
 returnn/frontend/audio/__init__.py,sha256=8mahwucBje8qHKw0bOvoySlvvD0rFKxviSvcAHSjiJY,67
 returnn/frontend/audio/mel.py,sha256=LNzC9aWWgLqua34bwxA--M9shtLlePfwLQ-HpvP2o54,7884
@@ -193,7 +193,7 @@ returnn/tf/frontend_low_level/__init__.py,sha256=34469k3KzMUIGowxReOZnbf6WdTjxY7
 returnn/tf/frontend_low_level/_backend.py,sha256=JwwRRIGnElqBC4bTImdB7w3U1u_SJESeZHYLmq86wog,24479
 returnn/tf/layers/__init__.py,sha256=Ngu-X84nWFgz7ndDu88DqoZ-5lUMMTQWH4g7N8pSoCg,72
 returnn/tf/layers/base.py,sha256=sUxEfh6WxaHWHG7O3cfxB6gG6YpEHkFKUJVayKvTBSI,152968
-returnn/tf/layers/basic.py,sha256=KOmEKBz5idk_-zL1XlKWmL4AIbgpwbHkGKZrszYrlM4,614963
+returnn/tf/layers/basic.py,sha256=zHDPLP97jSvYYZcMPqQVOVxFk6I1BfXd71XVfs0VIkQ,615386
 returnn/tf/layers/rec.py,sha256=3f6M_5aAMPvx7aAHdPV3VSFRHf7tjpp8lrXSzmk1I5c,548435
 returnn/tf/layers/segmental_model.py,sha256=wUyDZGr-eTVIIQWcsHLML0wtOxuWn_NFKOIrUKQcvoI,21515
 returnn/tf/layers/signal_processing.py,sha256=vRlkN7k7otk9_Qdv0qr_l6V0VT5Q6dO2MxwZWb2HH2M,52693
@@ -216,7 +216,7 @@ returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWI
 returnn/torch/data/returnn_dataset_wrapper.py,sha256=2CaDapzrlqahANuq-nyVAtv5ENHuM8A7okORwYJDisg,8006
 returnn/torch/data/tensor_utils.py,sha256=-Teqi--LLbt6q_5mDRdoHZHmPgSdC83W706ukif_YiU,1284
 returnn/torch/frontend/__init__.py,sha256=AA48HZnC17ASuKA0EWy8loZ-Bib_yUtqF4T1wYvjst4,62
-returnn/torch/frontend/_backend.py,sha256=L0SOhpm6c-9wbAUJIAgf9Qo8HVIyka2aHyAIJVt6W-A,101850
+returnn/torch/frontend/_backend.py,sha256=a9qcpUJrSDtH7KR6ZIpB4sijm6ztRlZ4myAe2P0dtaE,101875
 returnn/torch/frontend/_rand.py,sha256=1JgIkV2XmpgJD86zXZ-NCAe-QuoP2swr6NaS1oz3Qa8,1830
 returnn/torch/frontend/bridge.py,sha256=c_mVBCBo29sjm8Bhxarv00szwGPgxjwoIqAHOmceGQw,7842
 returnn/torch/frontend/raw_ops.py,sha256=lF0h-KtYYsdaaqQADylVZp9qzPskOOXA4MfmYDyx5IU,296
@@ -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.20250701.140328.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250701.140328.dist-info/METADATA,sha256=RByOzAa8sFb4gdGsaWiRUEOlZq3X9qqn4MEF0ziNPNc,5215
-returnn-1.20250701.140328.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20250701.140328.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250701.140328.dist-info/RECORD,,
+returnn-1.20250703.183400.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250703.183400.dist-info/METADATA,sha256=6b050edAnN_IixiTU7cKY9Y7GsV2u54Go4pMebtP1LA,5215
+returnn-1.20250703.183400.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20250703.183400.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250703.183400.dist-info/RECORD,,