returnn 1.20250701.140328__tar.gz → 1.20250703.183400__tar.gz

{returnn-1.20250701.140328/returnn.egg-info → returnn-1.20250703.183400}/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-1.20250703.183400/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250703.183400'
+long_version = '1.20250703.183400+git.a5c35a2'

{returnn-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/returnn/frontend/_native/module.hpp

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

{returnn-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/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-1.20250701.140328 → returnn-1.20250703.183400}/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 → returnn-1.20250703.183400/returnn.egg-info}/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-1.20250701.140328 → returnn-1.20250703.183400}/tests/test_rf_math.py

@@ -91,6 +91,24 @@ def test_squared_difference():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_log_add_exp():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "a": Tensor("a", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "b": Tensor("b", [in_dim], dtype="float32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_):
+        out = rf.log_add_exp(extern_data["a"], extern_data["b"])
+        out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
+
+    run_model(extern_data, lambda **_: rf.Module(), _forward_step)
+
+
 def test_abs_complex():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20250701.140328/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250701.140328'
-long_version = '1.20250701.140328+git.6646903'