returnn 1.20231213.133317__tar.gz → 1.20231214.122432__tar.gz

{returnn-1.20231213.133317/returnn.egg-info → returnn-1.20231214.122432}/PKG-INFO

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

returnn-1.20231214.122432/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20231214.122432'
+long_version = '1.20231214.122432+git.a043a50'

{returnn-1.20231213.133317 → returnn-1.20231214.122432}/returnn/frontend/attention.py

@@ -347,9 +347,11 @@ class RelPosSelfAttention(SelfAttentionBase):
     def __call__(self, source: Tensor, *, axis: Dim, **_kwargs) -> Tensor:
         """forward"""
         if self.learned_pos_emb is not None:
-            pos_emb, pos_emb_spatial_dim = self.learned_pos_emb(axis)
+            pos_emb, pos_emb_spatial_dim = self.learned_pos_emb(query_spatial_dim=axis, key_value_spatial_dim=axis)
         else:
-            pos_emb, pos_emb_spatial_dim = relative_positional_encoding(axis, self.pos_emb_feat_dim)
+            pos_emb, pos_emb_spatial_dim = relative_positional_encoding(
+                query_spatial_dim=axis, key_value_spatial_dim=axis, feat_dim=self.pos_emb_feat_dim
+            )
         if self.pos_emb_dropout:
             pos_emb = rf.dropout(pos_emb, self.pos_emb_dropout)
         if self.linear_pos is not None:
@@ -421,7 +423,7 @@ class LearnedRelativePositionalEncoding(rf.Module):
     https://github.com/rwth-i6/returnn_common/wiki/Relative-positional-encoding
     """
 
-    def __init__(self, feat_dim: Dim, *, clipping: int = 16, dtype: str = "float32"):
+    def __init__(self, feat_dim: Dim, *, clipping: int = 16, dtype: Optional[str] = None):
         """
         :param feat_dim: feature dim, for the emb matrix and output
         :param clipping: max distance to consider. emb matrix shape is [2 * clipping + 1, feat_dim].
@@ -434,7 +436,13 @@ class LearnedRelativePositionalEncoding(rf.Module):
         self.clipped_spatial_dim = Dim(2 * clipping + 1, name="learned-rel-pos")
         self.pos_emb = rf.Parameter((self.clipped_spatial_dim, self.feat_dim), dtype=dtype)
 
-    def __call__(self, spatial_dim: Dim) -> Tuple[Tensor, Dim]:
+    def __call__(
+        self,
+        *,
+        query_spatial_dim: Dim,
+        key_value_spatial_dim: Dim,
+        query_offset: Union[int, Tensor] = 0,
+    ) -> Tuple[Tensor, Dim]:
         """
         same interface as :func:`relative_positional_encoding`
 
@@ -442,9 +450,6 @@ class LearnedRelativePositionalEncoding(rf.Module):
           In the center is the rel pos i-j=0. All to the right are for i-j>0, all to the left for i-j<0.
         """
         # See also RelativePositionalEncodingLayer
-        query_offset = 0
-        query_spatial_dim = spatial_dim
-        key_value_spatial_dim = spatial_dim
         query_spatial_dim_m1 = query_spatial_dim - 1
 
         kv_pos_vec = rf.range_over_dim(key_value_spatial_dim)  # [kv_len]
@@ -464,14 +469,41 @@ class LearnedRelativePositionalEncoding(rf.Module):
         # Shift values to be >= 0. Each integer still uniquely identifies a relative position difference.
         indices = indices + self.clipping
 
-        encoding = rf.gather(self.pos_emb, indices=indices, axis=self.clipped_spatial_dim)  # [q_len,kv_len,n_out]
+        encoding = rf.gather(self.pos_emb, indices=indices, axis=self.clipped_spatial_dim)  # [out_spatial_dim,n_out]
         return encoding, out_spatial_dim
 
+    def full_matrix(
+        self,
+        *,
+        query_spatial_dim: Dim,
+        key_value_spatial_dim: Dim,
+        query_offset: Union[int, Tensor] = 0,
+    ) -> Tensor:
+        """
+        :return: as full matrix [query_spatial_dim,key_value_spatial_dim,feat_dim].
+            however, note that __call__ is usually to be preferred, as this gives a more efficient format.
+        """
+        # Very similar logic as in __call__.
+        kv_pos_vec = rf.range_over_dim(key_value_spatial_dim)  # [kv_len]
+        q_pos_vec = rf.range_over_dim(query_spatial_dim)  # [q_len]
+        indices = rf.combine_bc(kv_pos_vec, "-", q_pos_vec + query_offset)  # [q_len,kv_len]
+        indices = rf.clip_by_value(indices, -self.clipping, self.clipping)
+        indices = indices + self.clipping
+        encoding = rf.gather(self.pos_emb, indices=indices, axis=self.clipped_spatial_dim)  # [q_len,kv_len,n_out]
+        return encoding
+
 
 _relative_positional_encoding_cache = weakref.WeakKeyDictionary()  # run ctx -> (spatial_dim, feat_dim) -> enc
 
 
-def relative_positional_encoding(spatial_dim: Dim, feat_dim: Dim, *, dtype: str = None) -> Tuple[Tensor, Dim]:
+def relative_positional_encoding(
+    *,
+    query_spatial_dim: Dim,
+    key_value_spatial_dim: Dim,
+    feat_dim: Dim,
+    query_offset: int = 0,
+    dtype: Optional[str] = None,
+) -> Tuple[Tensor, Dim]:
     """
     Implements relative positional encoding, Transformer-XL style (https://arxiv.org/abs/1901.02860),
     as used for example by :class:`RelPosSelfAttention`.
@@ -494,25 +526,35 @@ def relative_positional_encoding(spatial_dim: Dim, feat_dim: Dim, *, dtype: str
     if not dtype:
         dtype = rf.get_default_float_dtype()
     cache = _relative_positional_encoding_cache.setdefault(rf.get_run_ctx(), {})
-    if (spatial_dim, feat_dim) in cache:
-        return cache[(spatial_dim, feat_dim)]
+    cache_key = (query_spatial_dim, key_value_spatial_dim, feat_dim, query_offset, dtype)
+    if cache_key in cache:
+        return cache[cache_key]
     import math
 
     with rf.control_flow_ctx(None):
-        position_pos = rf.range_over_dim(spatial_dim, dtype=dtype)
-        position_neg = -spatial_dim.get_dim_value_tensor() + rf.range_over_dim(spatial_dim - 1) + 1
-        position_neg = rf.cast(position_neg, dtype=dtype)
-        position, out_spatial_dim = rf.concat((position_neg, spatial_dim - 1), (position_pos, spatial_dim))
+        # See also RelativePositionalEncodingLayer, LearnedRelativePositionalEncoding
+        query_spatial_dim_m1 = query_spatial_dim - 1
+
+        kv_pos_vec = rf.range_over_dim(key_value_spatial_dim)  # [kv_len]
+        q_pos_vec = rf.range_over_dim(query_spatial_dim_m1)  # [q_len-1]
+
+        # See LearnedRelativePositionalEncoding.
+        indices, out_spatial_dim = rf.concat(
+            (q_pos_vec - query_spatial_dim_m1.get_dim_value_tensor(), query_spatial_dim_m1),
+            (kv_pos_vec, key_value_spatial_dim),
+        )
+        indices = indices - query_offset
+
         feat2_dim = feat_dim.div_left(2)
         div_term = rf.exp(rf.range_over_dim(feat2_dim, dtype=dtype) * -(2.0 * math.log(10000.0) / feat_dim.dimension))
-        arg_sin = rf.combine_bc(position, "*", div_term)
+        arg_sin = rf.combine_bc(rf.cast(indices, dtype), "*", div_term)
         arg_cos = arg_sin + math.pi / 2.0
         arg, feat_dim_ = rf.concat((arg_sin, feat2_dim), (arg_cos, feat2_dim))
         arg, feat_dim_ = rf.replace_dim(arg, in_dim=feat_dim_, out_dim=feat_dim)
         emb = rf.sin(arg)
         emb.verify_out_shape({out_spatial_dim, feat_dim}, allow_missing_implicit_dims=True)
         emb.feature_dim = feat_dim
-        cache[(spatial_dim, feat_dim)] = emb, out_spatial_dim
+        cache[cache_key] = emb, out_spatial_dim
         return emb, out_spatial_dim
 
 

{returnn-1.20231213.133317 → returnn-1.20231214.122432/returnn.egg-info}/PKG-INFO

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

{returnn-1.20231213.133317 → returnn-1.20231214.122432}/tests/test_rf_attention.py

@@ -81,7 +81,9 @@ def test_relative_positional_encoding():
 
     class _Net(rf.Module):
         def __call__(self, x: Tensor, *, axis: Dim) -> Tuple[Tensor, Dim]:
-            x, dim = rf.relative_positional_encoding(axis, in_dim)
+            x, dim = rf.relative_positional_encoding(
+                key_value_spatial_dim=axis, query_spatial_dim=axis, feat_dim=in_dim
+            )
             return x, dim
 
     # noinspection PyShadowingNames

returnn-1.20231213.133317/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20231213.133317'
-long_version = '1.20231213.133317+git.16419fa'