returnn 1.20240830.112737__tar.gz → 1.20240830.140746__tar.gz

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/PKG-INFO

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

returnn-1.20240830.140746/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240830.140746'
+long_version = '1.20240830.140746+git.d8709ff'

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/returnn/frontend/attention.py

@@ -16,6 +16,7 @@ __all__ = [
     "SelfAttention",
     "CausalSelfAttention",
     "CausalSelfAttentionState",
+    "RotaryPosSelfAttention",
     "RotaryPosCausalSelfAttention",
     "RelPosSelfAttention",
     "RelPosCausalSelfAttention",
@@ -264,6 +265,36 @@ class CausalSelfAttentionState(rf.State):
             self.accum_axis = accum_axis
 
 
+class RotaryPosSelfAttention(SelfAttention):
+    """
+    Rotary positional encoding (RoPE)-based self attention
+    """
+
+    def __call__(self, source: Tensor, *, axis: Dim) -> Tensor:
+        """forward"""
+        q, k, v = self.forward_qkv(source)
+
+        # Apply RoPE using sinusoidal positional encoding.
+        # Note: base is a bit different in rf.sinusoidal_positional_encoding (like the original)
+        # vs how it's commonly used for RoPE.
+        # log(base) / (dim / 2 - 1) = log(10_000) * 2 / dim
+        # <=> log(base) = log(10_000) * (dim / 2 - 1) * 2 / dim = log(10_000) * (1 - 2 / dim)
+        # <=> base = 10_000 ** (1 - 2 / dim)
+        pos_enc = rf.sinusoidal_positional_encoding(
+            spatial_dim=axis,
+            feat_dim=self.key_dim_per_head,
+            base=10_000 ** (1 - 2 / self.key_dim_per_head.dimension),
+        )  # [T,D]
+        q = _apply_rope(q, pos_enc, self.key_dim_per_head)
+        k = _apply_rope(k, pos_enc, self.key_dim_per_head)
+
+        kv_axis = Dim(None, name=f"{axis.name}-kv")
+        k, _ = rf.replace_dim(k, in_dim=axis, out_dim=kv_axis)
+        v, _ = rf.replace_dim(v, in_dim=axis, out_dim=kv_axis)
+        output = self.attention(q, k, v, kv_axis=kv_axis)
+        return output
+
+
 class RotaryPosCausalSelfAttention(CausalSelfAttention):
     """
     Rotary positional encoding (RoPE)-based causal self attention

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/returnn/frontend/container.py

@@ -66,14 +66,17 @@ class ModuleList(rf.Module, Generic[__ModT]):
         """module items"""
         return self._get_modules().items()
 
-    def __getitem__(self, idx: Union[slice, int]) -> Union[ModuleList[__ModT], __ModT]:
+    def __getitem__(self, idx: Union[slice, int, str]) -> Union[ModuleList[__ModT], __ModT]:
         if isinstance(idx, slice):
             return self.__class__(dict(list(self._get_modules().items())[idx]))
-        else:
+        elif isinstance(idx, str):
+            return getattr(self, idx)
+        elif isinstance(idx, int):
             return list(self._get_modules().values())[idx]
+        else:
+            raise TypeError(f"{self} __getitem__ Invalid idx type {type(idx).__name__}")
 
     def __setitem__(self, idx: Union[slice, int], module: Union[__ModT, Iterable[__ModT]]) -> None:
-        key = list(self._get_modules().keys())[idx]
         if isinstance(idx, slice):
             assert not idx.step or idx.step == 1  # not supported
             mod_items = list(self._get_modules().items())
@@ -95,8 +98,13 @@ class ModuleList(rf.Module, Generic[__ModT]):
                 assert not hasattr(self, k)
                 setattr(self, k, v)
                 i += 1
-        else:
+        elif isinstance(idx, str):
+            setattr(self, idx, module)
+        elif isinstance(idx, int):
+            key = list(self._get_modules().keys())[idx]
             setattr(self, key, _convert_to_module(module))
+        else:
+            raise TypeError(f"{self} __setitem___ Invalid idx type {type(idx).__name__}")
 
     def __delitem__(self, key: Union[slice, int]):
         if isinstance(key, slice):

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/returnn/frontend/encoder/conformer.py

@@ -295,6 +295,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         *,
         num_layers: int,
         input_layer: Optional[Union[ConformerConvSubsample, ISeqDownsamplingEncoder, rf.Module, Any]],
+        input_embedding_scale: float = 1.0,
         input_dropout: float = 0.1,
         ff_dim: Dim = NotSpecified,
         ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = NotSpecified,
@@ -312,6 +313,8 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         :param num_layers: the number of encoder layers
         :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))
         :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
@@ -338,6 +341,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
             if input_layer
             else None
         )
+        self.input_embedding_scale = input_embedding_scale
         self.input_dropout = input_dropout
 
         if not encoder_layer or isinstance(encoder_layer, (dict, type)):
@@ -389,6 +393,8 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         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)
         x = self.layers(x, spatial_dim=out_spatial_dim, collected_outputs=collected_outputs)
         return x, out_spatial_dim

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/returnn/frontend/encoder/e_branchformer.py

@@ -173,7 +173,13 @@ class FeedForwardConvGated(rf.Module):
     ):
         """
         :param out_dim: the encoder (e.g. E-Branchformer) model dim. (usually 256 or 512)
-        :param ff_dim: intermediate dimension. (usually 2048 or 3072, not necessarily factor 4 or 4*2/3)
+        :param ff_dim: intermediate dimension.
+            This is like cgmlp_linear_units/2 in ESPnet.
+            Note the 1/2 factor, which is because in ESPnet, you specify the total dimension,
+            before it is split for the gating,
+            while here, you specify the dimension for the gating part.
+            Common settings are 2048/2 or 3072/2.
+            In the paper, they mention a factor of 3 of the model dimension (factor 6 for ESPnet setting).
         :param kernel_size: for the depthwise convolution (usually 31)
         :param dropout:
         :param activation: activation function after the first linear layer, for both parts.
@@ -190,7 +196,7 @@ class FeedForwardConvGated(rf.Module):
         super().__init__()
 
         if ff_dim is NotSpecified:
-            ff_dim = out_dim * 6  # somewhat arbitrary. with 512, this is 3072.
+            ff_dim = out_dim * 3  # somewhat arbitrary. with 512, this is 3072/2.
         if isinstance(ff_dim, int):
             ff_dim = Dim(ff_dim, name="e-branchformer-ff-dim")
         if not isinstance(ff_dim, Dim):

{returnn-1.20240830.112737 → returnn-1.20240830.140746}/returnn.egg-info/PKG-INFO

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

returnn-1.20240830.112737/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240830.112737'
-long_version = '1.20240830.112737+git.9cd6289'