returnn 1.20250103.32004__tar.gz → 1.20250105.184816__tar.gz

{returnn-1.20250103.32004/returnn.egg-info → returnn-1.20250105.184816}/PKG-INFO

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

returnn-1.20250105.184816/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250105.184816'
+long_version = '1.20250105.184816+git.2c5e602'

{returnn-1.20250103.32004 → returnn-1.20250105.184816}/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.32004 → returnn-1.20250105.184816}/returnn/frontend/encoder/conformer.py

@@ -312,6 +312,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         sequential=rf.Sequential,
     ):
         """
+        :param in_dim: input features (e.g. MFCC)
         :param out_dim: the output feature dimension
         :param num_layers: the number of encoder layers
         :param input_layer: input/frontend/prenet with potential subsampling.

returnn-1.20250105.184816/returnn/frontend/encoder/conformer_v2.py

@@ -0,0 +1,169 @@
+"""
+Conformer model, variant of Transformer with additional convolution, introduced for speech recognition.
+Ref: https://arxiv.org/abs/2005.08100
+
+About details of the specific implementation and other implementations, see:
+https://github.com/rwth-i6/returnn_common/issues/233
+
+V2: Split frontend and main encoder.
+"""
+
+from __future__ import annotations
+from typing import Optional, Union, Any, Tuple, Dict, Callable
+import copy as _copy
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+from returnn.util.basic import NotSpecified
+from .base import ISeqDownsamplingEncoder, ISeqFramewiseEncoder
+from .conformer import ConformerConvSubsample, ConformerEncoderLayer
+
+
+class ConformerFrontend(ISeqDownsamplingEncoder):
+    """
+    This is just the combination of:
+    - input_layer (ConformerConvSubsample)
+    - input_projection (Linear without bias)
+    - input_embedding_scale
+    - input_dropout
+
+    This is intended to be used together with :class:`ConformerEncoderV2`.
+    """
+
+    def __init__(
+        self,
+        in_dim: Dim,
+        out_dim: Dim = Dim(512, name="conformer-enc-default-out-dim"),
+        *,
+        input_layer: Optional[Union[ConformerConvSubsample, ISeqDownsamplingEncoder, rf.Module, Any]],
+        input_embedding_scale: float = 1.0,
+        input_dropout: float = 0.1,
+    ):
+        """
+        :param in_dim: input features (e.g. MFCC)
+        :param out_dim: the output feature dimension
+        :param input_layer: input/frontend/prenet with potential subsampling.
+            (x, in_spatial_dim) -> (y, out_spatial_dim)
+        :param input_embedding_scale: applied after input_layer. 1.0 by default for historic reasons.
+            In std Transformer, also ESPnet E-Branchformer and Conformer, this is sqrt(out_dim).
+        :param input_dropout: applied after input_projection(input_layer(x))
+        """
+        super().__init__()
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.dropout_broadcast = rf.dropout_broadcast_default()
+
+        # TODO once we figured out good defaults, we would create ConformerConvSubsample here when not given
+        if callable(input_layer) or input_layer is None:
+            pass  # leave it as is
+        elif isinstance(input_layer, dict):
+            input_layer = rf.build_from_dict(input_layer, in_dim)
+            input_layer: ConformerConvSubsample  # maybe not true, but assume for some attribs
+        else:
+            raise TypeError(f"unexpected input_layer {input_layer!r}")
+        self.input_layer = input_layer
+        self.input_projection = (
+            rf.Linear(self.input_layer.out_dim if self.input_layer else self.in_dim, self.out_dim, with_bias=False)
+            if input_layer
+            else None
+        )
+        self.input_embedding_scale = input_embedding_scale
+        self.input_dropout = input_dropout
+
+    def __call__(self, source: Tensor, *, in_spatial_dim: Dim) -> Tuple[Tensor, Dim]:
+        """forward"""
+        if self.input_layer:
+            x_subsample, out_spatial_dim = self.input_layer(source, in_spatial_dim=in_spatial_dim)
+        else:
+            x_subsample, out_spatial_dim = source, in_spatial_dim
+        x = self.input_projection(x_subsample) if self.input_projection else x_subsample
+        if self.input_embedding_scale != 1.0:
+            x = x * self.input_embedding_scale
+        x = rf.dropout(x, self.input_dropout, axis=self.dropout_broadcast and self.out_dim)
+        return x, out_spatial_dim
+
+
+class ConformerEncoderV2(ISeqFramewiseEncoder):
+    """
+    Conformer encoder, without frontend.
+
+    V2: Without the input_layer / frontend module, i.e. just the conformer layers.
+    Use ConformerFrontend for the frontend.
+    To get the V1 case, add this in front:
+    """
+
+    def __init__(
+        self,
+        out_dim: Dim = Dim(512, name="conformer-enc-default-out-dim"),
+        *,
+        num_layers: int,
+        ff_dim: Dim = NotSpecified,
+        ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = NotSpecified,
+        dropout: float = 0.1,
+        conv_kernel_size: int = NotSpecified,
+        conv_norm: Union[rf.BatchNorm, type, Dict[str, Any], Any] = NotSpecified,
+        num_heads: int = 4,
+        att_dropout: float = 0.1,
+        encoder_layer: Optional[Union[ConformerEncoderLayer, rf.Module, type, Dict[str, Any], Any]] = None,
+        encoder_layer_opts: Optional[Dict[str, Any]] = None,
+        sequential=rf.Sequential,
+    ):
+        """
+        :param out_dim: the output feature dimension
+        :param num_layers: the number of encoder layers
+        :param ff_dim: the dimension of feed-forward layers. 2048 originally, or 4 times out_dim
+        :param ff_activation: activation function for feed-forward network
+        :param dropout: the dropout value for the FF block
+        :param conv_kernel_size: the kernel size of depthwise convolution in the conv block
+        :param conv_norm: used for the conv block. Batch norm originally
+        :param num_heads: the number of attention heads
+        :param att_dropout: attention dropout value
+        :param encoder_layer: an instance of :class:`ConformerEncoderLayer` or similar
+        :param encoder_layer_opts: options for the encoder layer
+        :param sequential:
+        """
+        super().__init__()
+
+        self.out_dim = out_dim
+        self.dropout = dropout
+        self.dropout_broadcast = rf.dropout_broadcast_default()
+
+        if not encoder_layer or isinstance(encoder_layer, (dict, type)):
+            encoder_layer_opts_ = dict(
+                out_dim=out_dim,
+                ff_dim=ff_dim,
+                ff_activation=ff_activation,
+                dropout=dropout,
+                conv_kernel_size=conv_kernel_size,
+                conv_norm=conv_norm,
+                num_heads=num_heads,
+                att_dropout=att_dropout,
+            )
+            encoder_layer_opts_ = {k: v for (k, v) in encoder_layer_opts_.items() if v is not NotSpecified}
+            if encoder_layer_opts:
+                encoder_layer_opts_.update(encoder_layer_opts)
+            if not encoder_layer:
+                encoder_layer = ConformerEncoderLayer(**encoder_layer_opts_)
+            elif isinstance(encoder_layer, type):
+                encoder_layer = encoder_layer(**encoder_layer_opts_)
+            elif isinstance(encoder_layer, dict):
+                # Note: Reuse all the encoder_layer_opts_.
+                # If this does not make sense for the specific encoder_layer class here,
+                # we would suggest to use a different ConformerEncoder class.
+                # (The alternative, to not reuse encoder_layer_opts_ here,
+                #  would probably be more confusing, as those options are all ignored then.
+                #  It's also not clear what args to pass then and what not.)
+                # (Maybe we should do a ConformerEncoderV2 if this is confusing here...)
+                encoder_layer_opts_ = {k: v for (k, v) in encoder_layer_opts_.items() if k not in encoder_layer}
+                encoder_layer = rf.build_from_dict(encoder_layer, **encoder_layer_opts_)
+            else:
+                raise TypeError(f"unexpected encoder_layer {encoder_layer!r}")
+        else:
+            if not callable(encoder_layer):
+                raise TypeError(f"{self}: invalid non-callable encoder_layer {encoder_layer!r}")
+
+        self.layers = sequential(_copy.deepcopy(encoder_layer) for _ in range(num_layers))
+
+    def __call__(self, x: Tensor, *, spatial_dim: Dim, collected_outputs: Optional[Dict[str, Tensor]] = None) -> Tensor:
+        """forward"""
+        return self.layers(x, spatial_dim=spatial_dim, collected_outputs=collected_outputs)

{returnn-1.20250103.32004 → returnn-1.20250105.184816}/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.32004 → returnn-1.20250105.184816}/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

{returnn-1.20250103.32004 → returnn-1.20250105.184816}/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.32004 → returnn-1.20250105.184816/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250103.32004
+Version: 1.20250105.184816
 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.32004 → returnn-1.20250105.184816}/returnn.egg-info/SOURCES.txt

@@ -218,6 +218,7 @@ returnn/frontend/decoder/transformer.py
 returnn/frontend/encoder/__init__.py
 returnn/frontend/encoder/base.py
 returnn/frontend/encoder/conformer.py
+returnn/frontend/encoder/conformer_v2.py
 returnn/frontend/encoder/e_branchformer.py
 returnn/frontend/encoder/transformer.py
 returnn/import_/__init__.py

returnn-1.20250103.32004/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250103.032004'
-long_version = '1.20250103.032004+git.fb9087a'