rasa-pro 3.9.18__py3-none-any.whl

rasa/utils/tensorflow/transformer.py

@@ -0,0 +1,640 @@
+from typing import Optional, Text, Tuple, Union
+
+import numpy as np
+import tensorflow as tf
+
+# TODO: The following is not (yet) available via tf.keras
+from keras.src.utils.control_flow_util import smart_cond
+from tensorflow.keras import backend as K
+
+import rasa.shared.utils.cli
+from rasa.utils.tensorflow.layers import RandomlyConnectedDense
+
+
+# from https://www.tensorflow.org/tutorials/text/transformer
+# and https://github.com/tensorflow/tensor2tensor
+class MultiHeadAttention(tf.keras.layers.Layer):
+    """Multi-headed attention layer.
+
+    Arguments:
+        units: Positive integer, output dim of hidden layer.
+        num_heads: Positive integer, number of heads
+            to repeat the same attention structure.
+        attention_dropout_rate: Float, dropout rate inside attention for training.
+        density: Approximate fraction of trainable weights (in
+            `RandomlyConnectedDense` layers).
+        unidirectional: Boolean, use a unidirectional or bidirectional encoder.
+        use_key_relative_position: Boolean, if 'True' use key
+            relative embeddings in attention.
+        use_value_relative_position: Boolean, if 'True' use value
+            relative embeddings in attention.
+        max_relative_position: Positive integer, max position for relative embeddings.
+        heads_share_relative_embedding: Boolean, if 'True'
+            heads will share relative embeddings.
+    """
+
+    def __init__(
+        self,
+        units: int,
+        num_heads: int,
+        attention_dropout_rate: float = 0.0,
+        density: float = 0.2,
+        unidirectional: bool = False,
+        use_key_relative_position: bool = False,
+        use_value_relative_position: bool = False,
+        max_relative_position: int = 5,
+        heads_share_relative_embedding: bool = False,
+    ) -> None:
+        super().__init__()
+
+        if units % num_heads != 0:
+            rasa.shared.utils.cli.print_error_and_exit(
+                f"Value Error: The given transformer size {units} should be a "
+                f"multiple of the number of attention heads {num_heads}."
+            )
+
+        self.num_heads = num_heads
+        self.units = units
+        self.attention_dropout_rate = attention_dropout_rate
+        self.unidirectional = unidirectional
+        self.use_key_relative_position = use_key_relative_position
+        self.use_value_relative_position = use_value_relative_position
+        self.relative_length = max_relative_position
+        self.relative_length += 1  # include current time
+        self.heads_share_relative_embedding = heads_share_relative_embedding
+
+        self._depth = units // self.num_heads
+
+        # process queries
+        self._query_dense_layer = RandomlyConnectedDense(
+            units=units, use_bias=False, density=density
+        )
+        # process keys
+        self._key_dense_layer = RandomlyConnectedDense(
+            units=units, use_bias=False, density=density
+        )
+        # process values
+        self._value_dense_layer = RandomlyConnectedDense(
+            units=units, use_bias=False, density=density
+        )
+        # process attention output
+        self._output_dense_layer = RandomlyConnectedDense(units=units, density=density)
+
+        self._create_relative_embeddings()
+
+    def _create_relative_embeddings(self) -> None:
+        """Create relative embeddings."""
+        relative_embedding_shape: Optional[
+            Union[Tuple[int, int], Tuple[int, int, int]]
+        ] = None
+        self.key_relative_embeddings = None
+        self.value_relative_embeddings = None
+
+        if self.use_key_relative_position or self.use_value_relative_position:
+            if not self.relative_length:
+                raise ValueError(
+                    f"Max relative position {self.relative_length} "
+                    f"should be > 0 when using relative attention."
+                )
+
+            if self.unidirectional:
+                relative_length = self.relative_length
+            else:
+                relative_length = 2 * self.relative_length - 1
+
+            if self.heads_share_relative_embedding:
+                relative_embedding_shape = (relative_length, self._depth)
+            else:
+                relative_embedding_shape = (
+                    self.num_heads,
+                    relative_length,
+                    self._depth,
+                )
+
+        if self.use_key_relative_position:
+            self.key_relative_embeddings = self.add_weight(
+                shape=relative_embedding_shape, name="key_relative_embeddings"
+            )
+
+        if self.use_value_relative_position:
+            self.value_relative_embeddings = self.add_weight(
+                shape=relative_embedding_shape, name="value_relative_embeddings"
+            )
+
+    def _pad_relative_embeddings(self, x: tf.Tensor, length: tf.Tensor) -> tf.Tensor:
+        # pad the left side to length
+        pad_left = x[:, :, :, :1, :]
+        pad_left = tf.tile(pad_left, (1, 1, 1, length - self.relative_length, 1))
+
+        # pad the right side to length
+        if self.unidirectional:
+            right_relative_length = 1  # current time
+            pad_right = tf.zeros_like(x[:, :, :, -1:, :])
+        else:
+            right_relative_length = self.relative_length
+            pad_right = x[:, :, :, -1:, :]
+        pad_right = tf.tile(pad_right, (1, 1, 1, length - right_relative_length, 1))
+
+        return tf.concat([pad_left, x, pad_right], axis=-2)
+
+    def _slice_relative_embeddings(self, x: tf.Tensor, length: tf.Tensor) -> tf.Tensor:
+        if self.unidirectional:
+            # pad the right side to relative_length
+            pad_right = tf.zeros_like(x[:, :, :, -1:, :])
+            pad_right = tf.tile(pad_right, (1, 1, 1, self.relative_length - 1, 1))
+            x = tf.concat([x, pad_right], axis=-2)
+
+        extra_length = self.relative_length - length
+        full_length = tf.shape(x)[-2]
+        return x[:, :, :, extra_length : full_length - extra_length, :]
+
+    def _relative_to_absolute_position(self, x: tf.Tensor) -> tf.Tensor:
+        """Universal method to convert tensor from relative to absolute indexing.
+
+        "Slides" relative embeddings by 45 degree.
+
+        Arguments:
+        x: A tensor of shape (batch, num_heads, length, relative_length, depth)
+            or (batch, num_heads, length, relative_length)
+
+        Returns:
+            A tensor of shape (batch, num_heads, length, length, depth)
+            or (batch, num_heads, length, length)
+        """
+        x_dim = len(x.shape)
+
+        if x_dim < 4 or x_dim > 5:
+            raise ValueError(
+                f"Relative tensor has a wrong shape {x.shape}, "
+                f"it should have 4 or 5 dimensions."
+            )
+        if x_dim == 4:
+            # add fake depth dimension
+            x = tf.expand_dims(x, axis=-1)
+
+        batch = tf.shape(x)[0]
+        num_heads = tf.shape(x)[1]
+        length = tf.shape(x)[2]
+        depth = tf.shape(x)[-1]
+
+        x = tf.cond(
+            length > self.relative_length,
+            lambda: self._pad_relative_embeddings(x, length),
+            lambda: self._slice_relative_embeddings(x, length),
+        )
+
+        # add a column of zeros to "slide" columns to diagonals through reshape
+        pad_shift = tf.zeros_like(x[:, :, :, -1:, :])
+        x = tf.concat([x, pad_shift], axis=-2)
+
+        # flatten length dimensions
+        x = tf.reshape(x, (batch, num_heads, -1, depth))
+        width = 2 * length
+
+        # add zeros so that the result of back reshape is still a matrix
+        pad_flat = tf.zeros_like(
+            x[:, :, : ((width - 1) - width * length % (width - 1)) % (width - 1), :]
+        )
+        x = tf.concat([x, pad_flat], axis=-2)
+
+        # "slide" columns to diagonals through reshape
+        x = tf.reshape(x, (batch, num_heads, -1, width - 1, depth))
+
+        # slice needed "diagonal" matrix
+        x = x[:, :, :-1, -length:, :]
+
+        if x_dim == 4:
+            # remove fake depth dimension
+            x = tf.squeeze(x, axis=-1)
+
+        return x
+
+    def _matmul_with_relative_keys(self, x: tf.Tensor) -> tf.Tensor:
+        y = self.key_relative_embeddings
+
+        if self.heads_share_relative_embedding:
+            matmul = tf.einsum("bhld,md->bhlm", x, y)
+        else:
+            matmul = tf.einsum("bhld,hmd->bhlm", x, y)
+
+        return self._relative_to_absolute_position(matmul)
+
+    def _tile_relative_embeddings(self, x: tf.Tensor, length: tf.Tensor) -> tf.Tensor:
+        if self.heads_share_relative_embedding:
+            x = tf.expand_dims(x, axis=0)  # add head dimension
+
+        x = tf.expand_dims(x, axis=1)  # add length dimension
+        x = tf.tile(x, (1, length, 1, 1))
+        return tf.expand_dims(x, axis=0)  # add batch dimension
+
+    def _squeeze_relative_embeddings(self, x: tf.Tensor) -> tf.Tensor:
+        x = tf.squeeze(x, axis=0)  # squeeze batch dimension
+        if self.heads_share_relative_embedding:
+            x = tf.squeeze(x, axis=1)  # squeeze head dimension
+        return x
+
+    def _matmul_with_relative_values(self, x: tf.Tensor) -> tf.Tensor:
+        y = self._tile_relative_embeddings(
+            self.value_relative_embeddings, tf.shape(x)[-2]
+        )
+        y = self._relative_to_absolute_position(y)
+        y = self._squeeze_relative_embeddings(y)
+
+        if self.heads_share_relative_embedding:
+            return tf.einsum("bhlm,lmd->bhld", x, y)
+        else:
+            return tf.einsum("bhlm,hlmd->bhld", x, y)
+
+    def _drop_attention_logits(
+        self, logits: tf.Tensor, pad_mask: tf.Tensor, training: tf.Tensor
+    ) -> tf.Tensor:
+        def droped_logits() -> tf.Tensor:
+            keep_prob = tf.random.uniform(tf.shape(logits), 0, 1) + pad_mask
+            drop_mask = tf.cast(
+                tf.less(keep_prob, self.attention_dropout_rate), logits.dtype
+            )
+
+            return logits + drop_mask * -1e9
+
+        return smart_cond(training, droped_logits, lambda: tf.identity(logits))
+
+    def _scaled_dot_product_attention(
+        self,
+        query: tf.Tensor,
+        key: tf.Tensor,
+        value: tf.Tensor,
+        pad_mask: tf.Tensor,
+        training: tf.Tensor,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Calculate the attention weights.
+
+        query, key, value must have matching leading dimensions.
+        key, value must have matching penultimate dimension,
+        i.e.: seq_len_k = seq_len_v.
+        The mask has different shapes depending on its type (padding or look ahead)
+        but it must be broadcastable for addition.
+
+        Arguments:
+            query: A tensor with shape (..., length, depth).
+            key: A tensor with shape (..., length, depth).
+            value: A tensor with shape (..., length, depth).
+            pad_mask: Float tensor with shape broadcastable
+                to (..., length, length). Defaults to None.
+
+        Returns:
+            output: A tensor with shape (..., length, depth).
+            attention_weights: A tensor with shape (..., length, length).
+        """
+        matmul_qk = tf.matmul(query, key, transpose_b=True)  # (..., length, length)
+
+        if self.use_key_relative_position:
+            matmul_qk += self._matmul_with_relative_keys(query)
+
+        # scale matmul_qk
+        dk = tf.cast(tf.shape(key)[-1], tf.float32)
+        logits = matmul_qk / tf.math.sqrt(dk)
+
+        # add the mask to the scaled tensor.
+        if pad_mask is not None:
+            logits += pad_mask * -1e9
+
+        # apply attention dropout before softmax to maintain attention_weights norm as 1
+        if self.attention_dropout_rate > 0:
+            logits = self._drop_attention_logits(logits, pad_mask, training)
+
+        # softmax is normalized on the last axis (length) so that the scores
+        # add up to 1.
+        attention_weights = tf.nn.softmax(logits, axis=-1)  # (..., length, length)
+
+        output = tf.matmul(attention_weights, value)  # (..., length, depth)
+        if self.use_value_relative_position:
+            output += self._matmul_with_relative_values(attention_weights)
+
+        return output, attention_weights
+
+    def _split_heads(self, x: tf.Tensor) -> tf.Tensor:
+        """Split the last dimension into (num_heads, depth).
+
+        Transpose the result such that the shape is
+        (batch_size, num_heads, length, depth)
+        """
+        x = tf.reshape(x, (tf.shape(x)[0], -1, self.num_heads, self._depth))
+        return tf.transpose(x, perm=[0, 2, 1, 3])
+
+    def _combine_heads(self, x: tf.Tensor) -> tf.Tensor:
+        """Inverse of split_heads.
+
+        Args:
+            x: A Tensor with shape [batch, num_heads, length, units / num_heads]
+
+        Returns:
+            A Tensor with shape [batch, length, units]
+        """
+        # (batch_size, length, num_heads, depth)
+        x = tf.transpose(x, perm=[0, 2, 1, 3])
+        # (batch_size, length, units)
+        return tf.reshape(x, (tf.shape(x)[0], -1, self.units))
+
+    # noinspection PyMethodOverriding
+    def call(
+        self,
+        query_input: tf.Tensor,
+        source_input: tf.Tensor,
+        pad_mask: Optional[tf.Tensor] = None,
+        training: Optional[Union[tf.Tensor, bool]] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Apply attention mechanism to query_input and source_input.
+
+        Arguments:
+            query_input: A tensor with shape [batch_size, length, input_size].
+            source_input: A tensor with shape [batch_size, length, input_size].
+            pad_mask: Float tensor with shape broadcastable
+                to (..., length, length). Defaults to None.
+            training: A bool, whether in training mode or not.
+
+        Returns:
+            Attention layer output with shape [batch_size, length, units]
+        """
+        if training is None:
+            training = K.learning_phase()
+
+        query = self._query_dense_layer(query_input)  # (batch_size, length, units)
+        key = self._key_dense_layer(source_input)  # (batch_size, length, units)
+        value = self._value_dense_layer(source_input)  # (batch_size, length, units)
+
+        query = self._split_heads(query)  # (batch_size, num_heads, length, depth)
+        key = self._split_heads(key)  # (batch_size, num_heads, length, depth)
+        value = self._split_heads(value)  # (batch_size, num_heads, length, depth)
+
+        attention, attention_weights = self._scaled_dot_product_attention(
+            query, key, value, pad_mask, training
+        )
+        # attention.shape == (batch_size, num_heads, length, depth)
+        # attention_weights.shape == (batch_size, num_heads, length, length)
+        attention = self._combine_heads(attention)  # (batch_size, length, units)
+
+        output = self._output_dense_layer(attention)  # (batch_size, length, units)
+
+        return output, attention_weights
+
+
+class TransformerEncoderLayer(tf.keras.layers.Layer):
+    """Transformer encoder layer.
+
+    The layer is composed of the sublayers:
+        1. Self-attention layer
+        2. Feed-forward network (which is 2 fully-connected layers)
+
+    Arguments:
+        units: Positive integer, output dim of hidden layer.
+        num_heads: Positive integer, number of heads
+            to repeat the same attention structure.
+        filter_units: Positive integer, output dim of the first ffn hidden layer.
+        dropout_rate: Float between 0 and 1; fraction of the input units to drop.
+        attention_dropout_rate: Float, dropout rate inside attention for training.
+        density: Fraction of trainable weights in `RandomlyConnectedDense` layers.
+        unidirectional: Boolean, use a unidirectional or bidirectional encoder.
+        use_key_relative_position: Boolean, if 'True' use key
+            relative embeddings in attention.
+        use_value_relative_position: Boolean, if 'True' use value
+            relative embeddings in attention.
+        max_relative_position: Positive integer, max position for relative embeddings.
+        heads_share_relative_embedding: Boolean, if 'True'
+            heads will share relative embeddings.
+    """
+
+    def __init__(
+        self,
+        units: int,
+        num_heads: int,
+        filter_units: int,
+        dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        density: float = 0.2,
+        unidirectional: bool = False,
+        use_key_relative_position: bool = False,
+        use_value_relative_position: bool = False,
+        max_relative_position: int = 5,
+        heads_share_relative_embedding: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self._layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self._mha = MultiHeadAttention(
+            units,
+            num_heads,
+            attention_dropout_rate,
+            density,
+            unidirectional,
+            use_key_relative_position,
+            use_value_relative_position,
+            max_relative_position,
+            heads_share_relative_embedding,
+        )
+        self._dropout = tf.keras.layers.Dropout(dropout_rate)
+
+        self._ffn_layers = [
+            tf.keras.layers.LayerNormalization(epsilon=1e-6),
+            RandomlyConnectedDense(
+                units=filter_units, activation=tf.nn.gelu, density=density
+            ),  # (batch_size, length, filter_units)
+            tf.keras.layers.Dropout(dropout_rate),
+            RandomlyConnectedDense(
+                units=units, density=density
+            ),  # (batch_size, length, units)
+            tf.keras.layers.Dropout(dropout_rate),
+        ]
+
+    def call(
+        self,
+        x: tf.Tensor,
+        pad_mask: Optional[tf.Tensor] = None,
+        training: Optional[Union[tf.Tensor, bool]] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Apply transformer encoder layer.
+
+        Arguments:
+            x: A tensor with shape [batch_size, length, units].
+            pad_mask: Float tensor with shape broadcastable
+                to (..., length, length). Defaults to None.
+            training: A bool, whether in training mode or not.
+
+        Returns:
+            Transformer encoder layer output with shape [batch_size, length, units]
+        """
+        if training is None:
+            training = K.learning_phase()
+
+        x_norm = self._layer_norm(x)  # (batch_size, length, units)
+        attn_out, attn_weights = self._mha(
+            x_norm, x_norm, pad_mask=pad_mask, training=training
+        )
+        attn_out = self._dropout(attn_out, training=training)
+        x += attn_out
+
+        ffn_out = x  # (batch_size, length, units)
+        for layer in self._ffn_layers:
+            ffn_out = layer(ffn_out, training=training)
+        x += ffn_out
+
+        # (batch_size, length, units), (batch_size, num_heads, length, length)
+        return x, attn_weights
+
+
+class TransformerEncoder(tf.keras.layers.Layer):
+    """Transformer encoder.
+
+    Encoder stack is made up of `num_layers` identical encoder layers.
+
+    Arguments:
+        num_layers: Positive integer, number of encoder layers.
+        units: Positive integer, output dim of hidden layer.
+        num_heads: Positive integer, number of heads
+            to repeat the same attention structure.
+        filter_units: Positive integer, output dim of the first ffn hidden layer.
+        reg_lambda: Float, regularization factor.
+        dropout_rate: Float between 0 and 1; fraction of the input units to drop.
+        attention_dropout_rate: Float, dropout rate inside attention for training.
+        density: Approximate fraction of trainable weights (in
+            `RandomlyConnectedDense` layers).
+        unidirectional: Boolean, use a unidirectional or bidirectional encoder.
+        use_key_relative_position: Boolean, if 'True' use key
+            relative embeddings in attention.
+        use_value_relative_position: Boolean, if 'True' use value
+            relative embeddings in attention.
+        max_relative_position: Positive integer, max position for relative embeddings.
+        heads_share_relative_embedding: Boolean, if 'True'
+            heads will share relative embeddings.
+        name: Optional name of the layer.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        units: int,
+        num_heads: int,
+        filter_units: int,
+        reg_lambda: float,
+        dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        density: float = 0.2,
+        unidirectional: bool = False,
+        use_key_relative_position: bool = False,
+        use_value_relative_position: bool = False,
+        max_relative_position: int = 5,
+        heads_share_relative_embedding: bool = False,
+        name: Optional[Text] = None,
+    ) -> None:
+        super().__init__(name=name)
+
+        self.units = units
+        self.unidirectional = unidirectional
+
+        l2_regularizer = tf.keras.regularizers.l2(reg_lambda)
+        self._embedding = RandomlyConnectedDense(
+            units=units, kernel_regularizer=l2_regularizer, density=density
+        )
+        # positional encoding helpers
+        self._angles = self._get_angles()
+        self._even_indices = np.arange(0, self.units, 2, dtype=np.int32)[:, np.newaxis]
+        self._odd_indices = np.arange(1, self.units, 2, dtype=np.int32)[:, np.newaxis]
+
+        self._dropout = tf.keras.layers.Dropout(dropout_rate)
+
+        self._enc_layers = [
+            TransformerEncoderLayer(
+                units,
+                num_heads,
+                filter_units,
+                dropout_rate,
+                attention_dropout_rate,
+                density,
+                unidirectional,
+                use_key_relative_position,
+                use_value_relative_position,
+                max_relative_position,
+                heads_share_relative_embedding,
+            )
+            for _ in range(num_layers)
+        ]
+        self._layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+
+    def _get_angles(self) -> np.ndarray:
+        array_2d = np.arange(self.units)[np.newaxis, :]
+        return 1 / np.power(10000, (2 * (array_2d // 2)) / np.float32(self.units))
+
+    def _positional_encoding(self, max_position: tf.Tensor) -> tf.Tensor:
+        max_position = tf.cast(max_position, dtype=tf.float32)
+        angle_rads = tf.range(max_position)[:, tf.newaxis] * self._angles
+
+        # transpose for easy slicing
+        angle_rads = tf.transpose(angle_rads, perm=[1, 0])
+        shape = tf.shape(angle_rads)
+        # apply sin to even indices in the array; 2i
+        sin_even = tf.sin(tf.gather_nd(angle_rads, self._even_indices))
+        pos_encoding_even = tf.scatter_nd(self._even_indices, sin_even, shape)
+        # apply cos to odd indices in the array; 2i+1
+        cos_odd = tf.cos(tf.gather_nd(angle_rads, self._odd_indices))
+        pos_encoding_odd = tf.scatter_nd(self._odd_indices, cos_odd, shape)
+        # combine even and odd positions and transpose back
+        pos_encoding = tf.transpose(pos_encoding_even + pos_encoding_odd, perm=[1, 0])
+        # add batch dimension
+        return tf.stop_gradient(pos_encoding[tf.newaxis, ...])
+
+    @staticmethod
+    def _look_ahead_pad_mask(max_position: tf.Tensor) -> tf.Tensor:
+        pad_mask = 1 - tf.linalg.band_part(tf.ones((max_position, max_position)), -1, 0)
+        return pad_mask[tf.newaxis, tf.newaxis, :, :]  # (1, 1, seq_len, seq_len)
+
+    def call(
+        self,
+        x: tf.Tensor,
+        pad_mask: Optional[tf.Tensor] = None,
+        training: Optional[Union[tf.Tensor, bool]] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Apply transformer encoder.
+
+        Arguments:
+            x: A tensor with shape [batch_size, length, input_size].
+            pad_mask: Float tensor with shape broadcastable
+                to (..., length, length). Defaults to None.
+            training: A bool, whether in training mode or not.
+
+        Returns:
+            Transformer encoder output with shape [batch_size, length, units]
+        """
+        # adding embedding and position encoding.
+        x = self._embedding(x)  # (batch_size, length, units)
+        x *= tf.math.sqrt(tf.cast(self.units, tf.float32))
+        x += self._positional_encoding(tf.shape(x)[1])
+        x = self._dropout(x, training=training)
+
+        if pad_mask is not None:
+            pad_mask = tf.squeeze(pad_mask, -1)  # (batch_size, length)
+            pad_mask = pad_mask[:, tf.newaxis, tf.newaxis, :]
+            # pad_mask.shape = (batch_size, 1, 1, length)
+            if self.unidirectional:
+                # add look ahead pad mask to emulate unidirectional behavior
+                pad_mask = tf.minimum(
+                    1.0, pad_mask + self._look_ahead_pad_mask(tf.shape(pad_mask)[-1])
+                )  # (batch_size, 1, length, length)
+
+        layer_attention_weights = []
+
+        for layer in self._enc_layers:
+            x, attn_weights = layer(x, pad_mask=pad_mask, training=training)
+            layer_attention_weights.append(attn_weights)
+
+        # if normalization is done in encoding layers, then it should also be done
+        # on the output, since the output can grow very large, being the sum of
+        # a whole stack of unnormalized layer outputs.
+        x = self._layer_norm(x)  # (batch_size, length, units)
+
+        # Keep the batch dimension on the first axis
+        attention_weights_as_output = tf.transpose(
+            tf.stack(layer_attention_weights), (1, 0, 2, 3, 4)
+        )
+
+        # (batch_size, length, units),
+        # (batch_size, num_layers, num_heads, length, length)
+        return x, attention_weights_as_output

rasa/utils/tensorflow/types.py

@@ -0,0 +1,6 @@
+from typing import Tuple, Union
+import tensorflow as tf
+import numpy as np
+
+BatchData = Union[Tuple[tf.Tensor, ...], Tuple[np.ndarray, ...]]
+MaybeNestedBatchData = Union[Tuple[BatchData, ...], BatchData]