returnn 1.20250103.10406__tar.gz → 1.20250104.234019__tar.gz

{returnn-1.20250103.10406/returnn.egg-info → returnn-1.20250104.234019}/PKG-INFO

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

returnn-1.20250104.234019/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250104.234019'
+long_version = '1.20250104.234019+git.62a72c8'

{returnn-1.20250103.10406 → returnn-1.20250104.234019}/returnn/frontend/dims.py

@@ -133,10 +133,13 @@ def dim_match_priority_when_needed(dim: Dim, *other_dims: Dim) -> Dim:
     return dim
 
 
-def num_elements_of_shape(dims: Union[Dim, Sequence[Dim]], *, use_mask: bool = True) -> Union[int, Tensor]:
+def num_elements_of_shape(
+    dims: Union[Dim, Sequence[Dim]], *, use_mask: bool = True, device: Optional[str] = None
+) -> Union[int, Tensor]:
     """
     :param dims:
     :param use_mask:
+    :param device: only for the case when we return a Tensor. by default, this is CPU (just as the size tensor).
     :return: num elements of a tensor of shape dims, properly considering masking
     """
     if isinstance(dims, Dim):
@@ -149,8 +152,10 @@ def num_elements_of_shape(dims: Union[Dim, Sequence[Dim]], *, use_mask: bool = T
     if all(dim.is_static() for dim in dims):
         n = 1
         for dim in dims:
-            n *= dim.dimension
+            n *= dim.size
         return n
+    if len(dims) == 1:
+        return dims[0].get_size_tensor(device=device)
 
     n: Union[int, Tensor] = 1
     postponed_dims = []
@@ -164,9 +169,9 @@ def num_elements_of_shape(dims: Union[Dim, Sequence[Dim]], *, use_mask: bool = T
                 related_dims.append(dim_)
         if not related_dims:
             if dim.is_static():
-                n *= dim.dimension
+                n *= dim.size
             else:
-                n *= dim.dyn_size_ext
+                n *= dim.get_size_tensor(device=device)
         else:
             postponed_dims.append(dim)
     if postponed_dims:

{returnn-1.20250103.10406 → returnn-1.20250104.234019}/returnn/frontend/math_.py

@@ -441,10 +441,15 @@ def log(a: Tensor) -> Tensor:
     return a._raw_backend.activation(a, "log")
 
 
+_SafeLogEps = {"float16": 6e-08, "bfloat16": 9.1835e-41, "float32": 1.4013e-45, "float64": 4.9407e-324}
+
+
 def safe_log(a: Tensor, *, eps: Optional[float] = None) -> Tensor:
     """safe_log"""
     if eps is None:
-        eps = {"float16": 6e-08, "bfloat16": 9.1835e-41, "float32": 1.4013e-45, "float64": 4.9407e-324}[a.dtype]
+        if a.dtype not in _SafeLogEps:
+            raise NotImplementedError(f"safe_log not implemented for dtype {a.dtype}")
+        eps = _SafeLogEps[a.dtype]
     # noinspection PyProtectedMember
     return a._raw_backend.safe_log(a, eps=eps)
 

{returnn-1.20250103.10406 → returnn-1.20250104.234019}/returnn/frontend/reduce.py

@@ -16,6 +16,7 @@ __all__ = [
     "reduce_min",
     "reduce_mean",
     "reduce_logsumexp",
+    "reduce_logmeanexp",
     "reduce_any",
     "reduce_all",
     "reduce_argmin",
@@ -106,6 +107,19 @@ def reduce_logsumexp(source: Tensor[T], *, axis: Union[Dim, Sequence[Dim]], use_
     return reduce(source=source, mode="logsumexp", axis=axis, use_mask=use_mask)
 
 
+def reduce_logmeanexp(source: Tensor[T], *, axis: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
+    """
+    Reduce the tensor along the given axis
+
+    :param source:
+    :param axis:
+    :param use_mask: if True (default), use the time mask (part of dim tag) to ignore padding frames
+    :return: tensor with axis removed
+    """
+    s = reduce_logsumexp(source, axis=axis, use_mask=use_mask)
+    return s - rf.safe_log(rf.cast(rf.num_elements_of_shape(axis, use_mask=use_mask, device=s.device), s.dtype))
+
+
 def reduce_any(source: Tensor[T], *, axis: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
     """
     Reduce the tensor along the given axis
@@ -203,6 +217,7 @@ class RunningMean(rf.Module):
         alpha: float,
         dtype: Optional[str] = None,
         is_prob_distribution: Optional[bool] = None,
+        update_only_in_train: bool = True,
     ):
         """
         :param in_dim: the dim of the mean vector, or the shape.
@@ -210,6 +225,8 @@ class RunningMean(rf.Module):
             Also called momentum. E.g. 0.1 is a common value, or less, like 0.001.
         :param dtype: the dtype of the mean vector
         :param is_prob_distribution: if True, will initialize the mean vector with 1/in_dim.
+        :param update_only_in_train: if True (default), will only update the mean vector in training mode.
+            False means it will always update.
         """
         super().__init__()
         self.in_dim = in_dim
@@ -221,15 +238,20 @@ class RunningMean(rf.Module):
         if is_prob_distribution:
             assert in_dim.dimension is not None
             self.mean.initial = 1.0 / in_dim.dimension
+        self.update_only_in_train = update_only_in_train
 
     def __call__(self, x: Tensor) -> Tensor:
         """
         :param x: shape [..., F]
         :return: shape [F]
         """
-        assert all(d in self.shape for d in x.dims)
-        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))
+
+        def _update_running_stats():
+            assert all(d in self.shape for d in x.dims)
+            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))
+
+        rf.cond((not self.update_only_in_train) or rf.get_run_ctx().train_flag, _update_running_stats, lambda: None)
         return self.mean
 
 

{returnn-1.20250103.10406 → returnn-1.20250104.234019}/returnn/tensor/_dim_extra.py

@@ -1910,6 +1910,7 @@ class _DimMixin:
 
     def get_size_tensor(self, *, device: Optional[str] = None) -> _t.Tensor:
         """
+        :param device: if None, will use CPU
         :return: size tensor, or dyn_size_ext if defined
         :rtype: _t.Tensor
         """

{returnn-1.20250103.10406 → returnn-1.20250104.234019/returnn.egg-info}/PKG-INFO

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

returnn-1.20250103.10406/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250103.010406'
-long_version = '1.20250103.010406+git.0dd1b3f'