Document AbsolutePositionalEncoding

tests/test_main.py

@@ -5,7 +5,7 @@
 import tensorflow as tf
 
 from transformerx.layers import (
-    MultiHeadAttention, PositionalEncoding, PositionWiseFFN, AddNorm,
+    MultiHeadAttention, AbsolutePositionalEncoding, PositionWiseFFN, AddNorm,
     TransformerEncoderBlock, TransformerEncoder, TransformerDecoderBlock, DotProductAttention,
 )
 from transformerx.txplot import Plot
@@ -25,7 +25,7 @@ def test_transpose_qkv():
 
 
 depth, num_steps = 32, 50
-pos_encoding = PositionalEncoding(depth, 0)
+pos_encoding = AbsolutePositionalEncoding(depth, 0)
 X = pos_encoding(tf.zeros((2, num_steps, depth)), training=False)
 P = pos_encoding.P[:, : X.shape[1], :]
 plotter = Plot()

transformerx/layers/__init__.py

@@ -1,7 +1,7 @@
 from .addnorm import AddNorm
 from .dot_product_attention import DotProductAttention
 from .multihead_attention import MultiHeadAttention
-from .positional_encoding import PositionalEncoding
+from .positional_encoding import AbsolutePositionalEncoding
 from .positionwise_ffn import PositionWiseFFN
 from .transformer_decoder import TransformerDecoder
 from .transformer_decoder_block import TransformerDecoderBlock

transformerx/layers/addnorm.py

@@ -94,29 +94,29 @@ def __init__(self, norm_shape: Tuple[int], dropout_rate: float = 0):
         self.dropout = tf.keras.layers.Dropout(dropout_rate)
         self.ln = tf.keras.layers.LayerNormalization(norm_shape)
 
-    def call(self, X: tf.Tensor, Y: tf.Tensor, **kwargs):
+    def call(self, x: tf.Tensor, residual: tf.Tensor, **kwargs):
         """Call AddNorm layer.
 
         Parameters
         ----------
-        X :
+        x :
             Input tensor
-        Y :
-            Input tensor 2
+        residual :
+            Residual input tensor
 
         Returns
         -------
         output :
             Added and normalized tensor
         """
-        if not isinstance(X, tf.Tensor):
+        if not isinstance(x, tf.Tensor):
             raise TypeError(
                 f"Expected a tensor for the "
-                f"argument 'X', but received: {X}"
+                f"argument 'x', but received: {x}"
             )
-        if not isinstance(Y, tf.Tensor):
+        if not isinstance(residual, tf.Tensor):
             raise TypeError(
                 f"Expected a tensor for the "
-                f"argument 'Y', but received: {Y}"
+                f"argument 'residual', but received: {residual}"
             )
-        return self.ln(self.dropout(Y, **kwargs) + X)
+        return self.ln(self.dropout(residual, **kwargs) + x)

transformerx/layers/multihead_attention.py

@@ -103,7 +103,7 @@ def __init__(
     def split_heads(self, X: tf.Tensor) -> tf.Tensor:
         """Transpose tensors for parallel computation of attention heads.
 
-        First transposition produces a tensor of shape X: (batch_size, num_heads, no. of queries or key-value pairs,
+        First transposition produces a tensor of shape x: (batch_size, num_heads, no. of queries or key-value pairs,
         depth / num_heads).
         Next it is rearranged to a new order (batch_size * num_heads, no. of queries or key-value pairs,
         depth / num_heads) which is then passed to the last rearrangement and returned.
@@ -114,19 +114,19 @@ def split_heads(self, X: tf.Tensor) -> tf.Tensor:
             The tensor to be transposed and prepared for the multi-head attention layer (i.e. queries, keys, and values)
         Returns
         -------
-        X : tf.Tensor
+        x : tf.Tensor
             Transposed tensor of shape ((batch_size * num_heads, no. of queries or key-value pairs, depth / num_heads)
         """
 
-        # X = tf.reshape(X, shape=(X.shape[0], X.shape[1], self.num_heads, -1))
+        # x = tf.reshape(x, shape=(x.shape[0], x.shape[1], self.num_heads, -1))
         X = rearrange(X, "b h (heads hidden) -> b h heads hidden", heads=self.num_heads)
-        # print("X reshaped: ", X.shape)
-        # X = tf.transpose(X, perm=(0, 2, 1, 3))
+        # print("x reshaped: ", x.shape)
+        # x = tf.transpose(x, perm=(0, 2, 1, 3))
         X = rearrange(X, "b d1 d2 d3 -> b d2 d1 d3")
-        # print("X transposed: ", X.shape)
-        # return tf.reshape(X, shape=(-1, X.shape[2], X.shape[3]))
+        # print("x transposed: ", x.shape)
+        # return tf.reshape(x, shape=(-1, x.shape[2], x.shape[3]))
         X = rearrange(X, "b d1 d2 d3 -> (b d1) d2 d3")
-        # print("X reshaped2: ", X.shape)
+        # print("x reshaped2: ", x.shape)
         return X
 
     def inverse_transpose_qkv(self, X):

transformerx/layers/positional_encoding.py

@@ -2,27 +2,70 @@
 import tensorflow as tf
 
 
-class PositionalEncoding(tf.keras.layers.Layer):
-    def __init__(self, num_hiddens, dropout, max_len=1000):
+class AbsolutePositionalEncoding(tf.keras.layers.Layer):
+    """Absolute positional encoding object.
+
+    Generate a sinusoid for each dimension of the positional encoding where wavelengths form a geometric progression
+    from :math:`2π` to :math:`(10000)2π`.
+
+    Notes
+    -----
+    Absolute Position Encodings are a type of position embeddings for [Transformer-based models] where positional
+    encodings are added to the input embeddings at the bottoms of the encoder and decoder stacks. The positional
+    encodings have the same dimension :math:`d_model` as the embeddings, so that the two can be summed. In the original
+    implementation, sine and cosine functions of different frequencies are used:
+
+    .. math::
+        PE(pos, 2i) = \sin(pos^{2i/d_{model}})
+
+        PE(pos, 2i+1) = \cos(pos^{2i/d_{model}})
+
+    where  is the position and  is the dimension.
+
+
+    Parameters
+    ----------
+    depth :
+        Length of the positional encoding's hidden units.
+    dropout_rate :
+        Float between 0 and 1. Fraction of the input units to drop.
+    max_len :
+        Maximum length of the steps to calculate sinusoid
+
+    Examples
+    --------
+    >>> depth, num_steps = 32, 50
+    >>> pos_encoding = AbsolutePositionalEncoding(depth, dropout_rate=0.2)
+    >>> x = tf.zeros((1, num_steps, depth))
+    >>> print(x.shape)
+    (1, 50, 32)
+    >>> X = pos_encoding(x, training=False)
+    >>> P = pos_encoding.P[:, : X.shape[1], :]
+    >>> print(X.shape)
+    (1, 50, 32)
+    >>> print(P.shape)
+    (1, 50, 32)
+    """
+
+    def __init__(self, depth, dropout_rate=0, max_len=1000):
         super().__init__()
-        self.dropout = tf.keras.layers.Dropout(dropout)
+        self.dropout = tf.keras.layers.Dropout(dropout_rate)
         # Create a long enough P
 
-        self.P = np.zeros((1, max_len, num_hiddens))
-        print("P.shape", self.P.shape)
+        self.P = np.zeros((1, max_len, depth))
         X = np.arange(max_len, dtype=np.float32).reshape(-1, 1) / np.power(
-                10000, np.arange(0, num_hiddens, 2, dtype=np.float32) / num_hiddens
+            10000, np.arange(0, depth, 2, dtype=np.float32) / depth
         )
 
         self.P[:, :, 0::2] = tf.sin(
-                X
+            X
         )  # x[low::stride] -> positions: 0, 2, 4, ... of all rows and columns
         self.P[:, :, 1::2] = tf.cos(
-                X
+            X
         )  # x[low::stride] -> positions: 1, 3, 5 , ... of all rows and columns
 
     def call(self, X, **kwargs):
-        # print("X.shape[1]: ", X.shape[1])
-        # print("self.P[:, : X.shape[1], :]: ", self.P[:, : X.shape[1], :].shape)
+        # print("x.shape[1]: ", x.shape[1])
+        # print("self.P[:, : x.shape[1], :]: ", self.P[:, : x.shape[1], :].shape)
         X = X + self.P[:, : X.shape[1], :]
         return self.dropout(X, **kwargs)

transformerx/layers/transformer_decoder.py

@@ -1,6 +1,6 @@
 import tensorflow as tf
 
-from transformerx.layers.positional_encoding import PositionalEncoding
+from transformerx.layers.positional_encoding import AbsolutePositionalEncoding
 from transformerx.layers.transformer_decoder_block import TransformerDecoderBlock
 
 
@@ -21,7 +21,7 @@ def __init__(
         self.depth = depth
         self.n_blocks = n_blocks
         self.embedding = tf.keras.layers.Embedding(vocab_size, depth)
-        self.pos_encoding = PositionalEncoding(depth, dropout)
+        self.pos_encoding = AbsolutePositionalEncoding(depth, dropout)
         self.blocks = [
             TransformerDecoderBlock(
                     depth,

transformerx/layers/transformer_encoder.py

@@ -1,6 +1,6 @@
 import tensorflow as tf
 
-from transformerx.layers.positional_encoding import PositionalEncoding
+from transformerx.layers.positional_encoding import AbsolutePositionalEncoding
 from transformerx.layers.transformer_encoder_block import TransformerEncoderBlock
 
 
@@ -22,7 +22,7 @@ def __init__(
         self.depth = depth
         self.n_blocks = n_blocks
         self.embedding = tf.keras.layers.Embedding(vocab_size, depth)
-        self.pos_encoding = PositionalEncoding(depth, dropout)
+        self.pos_encoding = AbsolutePositionalEncoding(depth, dropout)
         self.blocks = [
             TransformerEncoderBlock(
                     depth,

transformerx/utils.py

@@ -6,7 +6,7 @@
 def masked_softmax(X, valid_lens):
     """Perform softmax operation by masking elements on the last axis."""
 
-    # X: 3D tensor, valid_lens: 1D or 2D tensor
+    # x: 3D tensor, valid_lens: 1D or 2D tensor
     def _sequence_mask(X, valid_len, value=0):
         maxlen = X.shape[1]
         mask = tf.range(start=0, limit=maxlen, dtype=tf.float32)[None, :] < tf.cast(