mttf 1.3.6__py3-none-any.whl

mt/keras_src/layers_src/transformer_encoder.py

@@ -0,0 +1,246 @@
+from .. import layers, activations, initializers
+
+
+class MTTransformerEncoder(layers.Layer):
+    """Transformer encoder.
+
+    This layer was ripped from [KerasHub](https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/layers/modeling/transformer_encoder.py)
+    by MT. It has been adapted and simplified to work with Keras 2. It does not
+    take as input attention mask or padding mask, but instead computes them
+    internally.
+
+    This class follows the architecture of the transformer encoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up an encoder.
+
+    Args:
+        intermediate_dim: int, the hidden size of feedforward network.
+        num_heads: int, the number of heads in the
+            `keras.layers.MultiHeadAttention` layer.
+        dropout: float. the dropout value, shared by
+            `keras.layers.MultiHeadAttention` and feedforward network.
+            Defaults to `0.`.
+        activation: string or `keras.activations`. the
+            activation function of feedforward network.
+            Defaults to `"relu"`.
+        layer_norm_epsilon: float. The epsilon value in layer
+            normalization components. Defaults to `1e-5`.
+        kernel_initializer: string or `keras.initializers` initializer.
+            The kernel initializer for the dense and multiheaded
+            attention layers. Defaults to `"glorot_uniform"`.
+        bias_initializer: string or `keras.initializers` initializer.
+            The bias initializer for the dense and multiheaded
+            attention layers. Defaults to `"zeros"`.
+        normalize_first: bool. If True, the inputs to the
+            attention layer and the intermediate dense layer  are normalized
+            (similar to GPT-2). If set to False, outputs of attention layer and
+            intermediate dense layer are normalized (similar to BERT).
+            Defaults to `False`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Example:
+
+    ```python
+    # Create a single transformer encoder layer.
+    encoder = keras_hub.layers.MTTransformerEncoder(
+        intermediate_dim=64, num_heads=8)
+
+    # Create a simple model containing the encoder.
+    input = keras.Input(shape=(10, 64))
+    output = encoder(input)
+    model = keras.Model(inputs=input, outputs=output)
+
+    # Call encoder on the inputs.
+    input_data = np.random.uniform(size=(2, 10, 64))
+    output = model(input_data)
+    ```
+
+    References:
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
+    """
+
+    def __init__(
+        self,
+        intermediate_dim,
+        num_heads,
+        dropout=0,
+        activation="relu",
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        normalize_first=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.intermediate_dim = intermediate_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.activation = activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.bias_initializer = initializers.get(bias_initializer)
+        self.normalize_first = normalize_first
+        self.supports_masking = True
+
+    def build(self, inputs_shape):
+        # Infer the dimension of our hidden feature size from the build shape.
+        hidden_dim = inputs_shape[-1]
+        # Attention head size is `hidden_dim` over the number of heads.
+        key_dim = int(hidden_dim // self.num_heads)
+        if key_dim == 0:
+            raise ValueError(
+                "Attention `key_dim` computed cannot be zero. "
+                f"The `hidden_dim` value of {hidden_dim} has to be equal to "
+                f"or greater than `num_heads` value of {self.num_heads}."
+            )
+
+        # Self attention layers.
+        self._self_attention_layer = layers.MultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=key_dim,
+            dropout=self.dropout,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
+            name="self_attention_layer",
+        )
+        if hasattr(self._self_attention_layer, "_build_from_signature"):
+            self._self_attention_layer._build_from_signature(
+                query=inputs_shape,
+                value=inputs_shape,
+            )
+        else:
+            self._self_attention_layer.build(
+                query_shape=inputs_shape,
+                value_shape=inputs_shape,
+            )
+        self._self_attention_layer_norm = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="self_attention_layer_norm",
+        )
+        self._self_attention_layer_norm.build(inputs_shape)
+        self._self_attention_dropout = layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
+        )
+
+        # Feedforward layers.
+        self._feedforward_layer_norm = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
+        )
+        self._feedforward_layer_norm.build(inputs_shape)
+        self._feedforward_intermediate_dense = layers.Dense(
+            self.intermediate_dim,
+            activation=self.activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
+            name="feedforward_intermediate_dense",
+        )
+        self._feedforward_intermediate_dense.build(inputs_shape)
+        self._feedforward_output_dense = layers.Dense(
+            hidden_dim,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
+        )
+        intermediate_shape = list(inputs_shape)
+        intermediate_shape[-1] = self.intermediate_dim
+        self._feedforward_output_dense.build(tuple(intermediate_shape))
+        self._feedforward_dropout = layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
+        )
+        self.built = True
+
+    def call(
+        self,
+        inputs,
+        training=None,
+        return_attention_scores=False,
+    ):
+        """Forward pass of the MTTransformerEncoder.
+
+        Args:
+            inputs: a Tensor. The input data to MTTransformerEncoder, should be
+                of shape [batch_size, sequence_length, hidden_dim].
+            training: a boolean indicating whether the layer should behave in
+                training mode or in inference mode.
+            return_attention_scores: a boolean indicating whether the output
+                should be `(attention_output, attention_scores)` if `True` or
+                `attention_output` if `False`. Defaults to `False`.
+
+        Returns:
+            A Tensor of the same shape as the `inputs`.
+        """
+        x = inputs  # Intermediate result.
+
+        # Self attention block.
+        residual = x
+        if self.normalize_first:
+            x = self._self_attention_layer_norm(x)
+
+        if return_attention_scores:
+            x, attention_scores = self._self_attention_layer(
+                query=x,
+                value=x,
+                return_attention_scores=return_attention_scores,
+                training=training,
+            )
+        else:
+            x = self._self_attention_layer(
+                query=x,
+                value=x,
+                training=training,
+            )
+
+        x = self._self_attention_dropout(x, training=training)
+        x = x + residual
+        if not self.normalize_first:
+            x = self._self_attention_layer_norm(x)
+
+        # Feedforward block.
+        residual = x
+        if self.normalize_first:
+            x = self._feedforward_layer_norm(x)
+        x = self._feedforward_intermediate_dense(x)
+        x = self._feedforward_output_dense(x)
+        x = self._feedforward_dropout(x, training=training)
+        x = x + residual
+        if not self.normalize_first:
+            x = self._feedforward_layer_norm(x)
+
+        if return_attention_scores:
+            return x, attention_scores
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "dropout": self.dropout,
+                "activation": activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": initializers.serialize(
+                    self.bias_initializer
+                ),
+                "normalize_first": self.normalize_first,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, inputs_shape):
+        return inputs_shape

mt/keras_src/layers_src/utils.py

@@ -0,0 +1,88 @@
+"""Useful subroutines dealing with GPU devices."""
+
+from mt import tp, tfc
+
+
+def conv2d(name_scope: tfc.NameScope, x, filters, kernel_size, **kwargs):
+    """Wrapper of Keras Conv2D layer with a LayerNormalization layer.
+
+    Parameters
+    ----------
+    name_scope : mt.tfc.NameScope
+        the name scope. For every conv2d invocation, the name scope is iterated.
+    x : tensor-like
+        Keras tensor or TF tensor as input
+    filters : int
+        The dimensionality of the output space (i.e. the number of output filters in the
+        convolution).
+    kernel_size : int or tuple or list
+        An integer or tuple/list of 2 integers, specifying the height and width of the 2D
+        convolution window. Can be a single integer to specify the same value for all spatial
+        dimensions.
+    **kwargs : dict
+        all other keyword arguments to be passed as-is to Conv2D layer construction
+
+    Returns
+    -------
+    tensor-like
+        TF tensor as output
+    """
+
+    from .. import layers
+
+    next(name_scope)
+    x = layers.LayerNormalization(name=name_scope("prenorm"))(x)
+    x = layers.Conv2D(filters, kernel_size, name=name_scope("conv"), **kwargs)(x)
+
+    return x
+
+
+def dense2d(
+    name_scope: tfc.NameScope, x, filters, kernel_size, activation="tanh", **kwargs
+):
+    """Wrapper of Keras Conv2D layer with a LayerNormalization layer.
+
+    TBD. But basically, prenorm, then expand to twice the dim, then prenorm, then project to target
+    dim with target kernel size
+
+    Parameters
+    ----------
+    name_scope : mt.tfc.NameScope
+        the name scope. For every conv2d invocation, the name scope is iterated.
+    x : tensor-like
+        Keras tensor or TF tensor as input
+    filters : int
+        The dimensionality of the output space (i.e. the number of output filters in the
+        convolution).
+    kernel_size : int or tuple or list
+        An integer or tuple/list of 2 integers, specifying the height and width of the 2D
+        convolution window. Can be a single integer to specify the same value for all spatial
+        dimensions.
+    activation : object
+        the activation of the last conv layer
+    **kwargs : dict
+        all other keyword arguments to be passed as-is to Conv2D layer construction
+
+    Returns
+    -------
+    tensor-like
+        TF tensor as output
+    """
+
+    from .. import layers
+
+    next(name_scope)
+    x = layers.LayerNormalization(name=name_scope("expand_prenorm"))(x)
+    x = layers.Conv2D(
+        x.shape[3] * 2, 1, name=name_scope("expand"), activation="relu", **kwargs
+    )(x)
+    x = layers.LayerNormalization(name=name_scope("project_prenorm"))(x)
+    x = layers.Conv2D(
+        filters,
+        kernel_size,
+        name=name_scope("project"),
+        activation=activation,
+        **kwargs
+    )(x)
+
+    return x

mt/keras_src/layers_src/var_regularizer.py

@@ -0,0 +1,38 @@
+import tensorflow as tf
+
+from mt import tp
+
+from .. import layers
+
+
+class VarianceRegularizer(layers.Layer):
+    """A regularizer on the variance of the input tensor.
+
+    Negative rate for making the variance larger. Positive rate for making the variance smaller.
+    """
+
+    def __init__(self, rate=1e-2, l_axes: list = [-1], **kwargs):
+        super(VarianceRegularizer, self).__init__(**kwargs)
+        self.rate = rate
+        self.l_axes = l_axes
+
+    def call(self, x):
+        mean = tf.reduce_mean(x, axis=self.l_axes, keepdims=True)
+        err = x - mean
+        esq = err * err
+        var = tf.reduce_mean(esq, axis=self.l_axes)
+        sum_var = tf.reduce_sum(var)
+        self.add_loss(self.rate * sum_var)
+        return x
+
+    call.__doc__ = layers.Layer.call.__doc__
+
+    def get_config(self):
+        config = {
+            "rate": self.rate,
+            "l_axes": self.l_axes,
+        }
+        base_config = super(VarianceRegularizer, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    get_config.__doc__ = layers.Layer.get_config.__doc__

mt/tf/__init__.py

@@ -0,0 +1,10 @@
+from .mttf_version import version as __mttf_version__
+from .init import init as mttf_init
+from .utils import *
+from tensorflow import *
+import tensorflow.compiler as compiler
+import tensorflow.core as core
+import tensorflow.keras as keras
+import tensorflow.lite as lite
+import tensorflow.python as python
+import tensorflow.tools as tools

mt/tf/init.py

@@ -0,0 +1,25 @@
+"""Initialises TensorFlow, monkey-patching if necessary."""
+
+from packaging.version import Version
+
+__all__ = ["init"]
+
+
+def init():
+    """Initialises tensorflow, monkey-patching if necessary."""
+
+    import tensorflow
+    import sys
+
+    tf_ver = Version(tensorflow.__version__)
+
+    if tf_ver < Version("2.10"):
+        raise ImportError(
+            f"The minimum TF version that mttf supports is 2.10. Your TF is {tf_ver}. "
+            "Please upgrade."
+        )
+
+    return tensorflow
+
+
+init()

mt/tf/keras_applications/__init__.py

@@ -0,0 +1,5 @@
+from mt.logg import logger
+
+logger.warn_module_move("mt.tf.keras_applications", "mt.keras.applications")
+
+from mt.keras.applications import *

mt/tf/keras_layers/__init__.py

@@ -0,0 +1,5 @@
+from mt.logg import logger
+
+logger.warn_module_move("mt.tf.keras_layers", "mt.keras.layers")
+
+from mt.keras.layers import *

mt/tf/mttf_version.py

@@ -0,0 +1,5 @@
+MAJOR_VERSION = 1
+MINOR_VERSION = 0
+PATCH_VERSION = 0
+version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
+__all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

mt/tf/utils.py

@@ -0,0 +1,44 @@
+"""Useful subroutines dealing with GPU devices."""
+
+from mt import tp, np, tfc
+
+__api__ = ["gpus_in_tf_format", "as_floatx", "sigmoid", "asigmoid"]
+
+
+NameScope = tfc.NameScope  # for backward compatibility
+
+
+def gpus_in_tf_format(gpus):
+    """Converts a gpu list or a gpu count into a list of GPUs in TF format."""
+
+    if isinstance(gpus, int):
+        gpus = range(gpus)
+    return ["/GPU:{}".format(x) for x in gpus]
+
+
+def as_floatx(x):
+    """Ensures that a tensor is of dtype floatx."""
+
+    import tensorflow as tf
+
+    if not np.issubdtype(x.dtype.as_numpy_dtype, np.floating):
+        x = tf.cast(x, tf.keras.backend.floatx())
+    return x
+
+
+def sigmoid(x):
+    """Stable sigmoid, taken from tfp."""
+
+    import tensorflow as tf
+
+    x = tf.convert_to_tensor(x)
+    cutoff = -20 if x.dtype == tf.float64 else -9
+    return tf.where(x < cutoff, tf.exp(x), tf.math.sigmoid(x))
+
+
+def asigmoid(y):
+    """Inverse of sigmoid, taken from tfp."""
+
+    import tensorflow as tf
+
+    return tf.math.log(y) - tf.math.log1p(-y)

mt/tf/version.py

@@ -0,0 +1,5 @@
+MAJOR_VERSION = 1
+MINOR_VERSION = 3
+PATCH_VERSION = 6
+version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
+__all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']