mttf 1.1.5__py3-none-any.whl → 1.1.9__py3-none-any.whl

mt/keras/__init__.py

@@ -0,0 +1 @@
+"""MT package that represents the working Keras 2 from the system."""

mt/keras/base/__init__.py

@@ -0,0 +1,10 @@
+from .base import keras_version, keras_source
+
+if keras_source == "tf_keras":
+    from tf_keras import *
+elif keras_source == "keras":
+    from keras import *
+elif keras_source == "tensorflow.keras":
+    from tensorflow.keras import *
+else:
+    raise ImportError(f"Unknown value '{keras_source}' for variable 'keras_source'.")

mt/keras/base/base.py

@@ -0,0 +1,28 @@
+"""Determines the working Keras 2 from the system to be used by mt.keras."""
+
+from packaging.version import Version
+import tensorflow as tf
+
+tf_ver = Version(tf.__version__)
+if tf_ver >= Version("2.16"):
+    try:
+        import tf_keras
+    except:
+        raise ImportError(
+            f"mt.keras can only work with Keras 2. You have TF version {tf_ver}. Please install tf_keras."
+        )
+    keras_version = tf_keras.__version__
+    keras_source = "tf_keras"
+else:
+    try:
+        import keras
+
+        kr_ver = Version(keras.__version__)
+    except ImportError:
+        kr_ver = None
+    if kr_ver is None or kr_ver >= Version("3.0"):
+        keras_version = tf.__version__
+        keras_source = "tensorflow.python"
+    else:
+        keras_version = keras.__version__
+        keras_source = "keras"

mt/keras/layers/__init__.py

@@ -0,0 +1,39 @@
+from ..base import layers as _layers
+
+for _x, _y in _layers.__dict__.items():
+    if x.startswith("_"):
+        continue
+    globals()[_x] = _y
+globals.__doc__ = _layers.__doc__
+
+from .identical import *
+from .floor import *
+from .var_regularizer import *
+from .simple_mha import *
+from .image_sizing import *
+from .counter import Counter
+from .normed_conv2d import NormedConv2D
+from .utils import *
+
+
+__api__ = [
+    "Identical",
+    "Floor",
+    "VarianceRegularizer",
+    "SimpleMHA2D",
+    "MHAPool2D",
+    "DUCLayer",
+    "Downsize2D",
+    "Upsize2D",
+    "Downsize2D_V2",
+    "Upsize2D_V2",
+    "Downsize2D_V3",
+    "Downsize2D_V4",
+    "DownsizeX2D",
+    "UpsizeX2D",
+    "DownsizeY2D",
+    "UpsizeY2D",
+    "Counter",
+    "conv2d",
+    "dense2d",
+]

mt/{tf/keras_layers → keras/layers}/counter.py

@@ -1,12 +1,12 @@
 import tensorflow as tf
-import tensorflow.keras as tk
+from ..base import layers, initializers
 
 
-class Counter(tk.layers.Layer):
+class Counter(layers.Layer):
     """A layer that counts from 0 during training and does nothing during inference."""
 
     def build(self, input_shape):
-        initializer = tk.initializers.Constant(value=0.0)
+        initializer = initializers.Constant(value=0.0)
         self.counter = self.add_weight(
             name="counter", shape=(1,), initializer=initializer
         )
@@ -19,9 +19,9 @@ class Counter(tk.layers.Layer):
         y = tf.stop_gradient(y) * 0.0
         return self.counter + y
 
-    call.__doc__ = tk.layers.Layer.call.__doc__
+    call.__doc__ = layers.Layer.call.__doc__
 
     def compute_output_shape(self, input_shape):
         return (1,)
 
-    compute_output_shape.__doc__ = tk.layers.Layer.compute_output_shape.__doc__
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__

mt/{tf/keras_layers → keras/layers}/floor.py

@@ -1,4 +1,5 @@
 import tensorflow as tf
+from ..base import layers
 
 
 @tf.custom_gradient
@@ -9,15 +10,15 @@ def floor(x):
     return tf.math.floor(x), grad
 
 
-class Floor(tf.keras.layers.Layer):
+class Floor(layers.Layer):
     """TensorFlow floor but gradient is identity."""
 
     def call(self, x):
         return floor(x)
 
-    call.__doc__ = tf.keras.layers.Layer.call.__doc__
+    call.__doc__ = layers.Layer.call.__doc__
 
     def compute_output_shape(self, input_shape):
         return input_shape
 
-    compute_output_shape.__doc__ = tf.keras.layers.Layer.compute_output_shape.__doc__
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__

mt/{tf/keras_layers → keras/layers}/identical.py

@@ -1,15 +1,15 @@
-import tensorflow.keras.layers as _kl
+from ..base import layers
 
 
-class Identical(_kl.Layer):
+class Identical(layers.Layer):
     """An identical layer, mainly for renaming purposes."""
 
     def call(self, x):
         return x
 
-    call.__doc__ = _kl.Layer.call.__doc__
+    call.__doc__ = layers.Layer.call.__doc__
 
     def compute_output_shape(self, input_shape):
         return input_shape
 
-    compute_output_shape.__doc__ = _kl.Layer.compute_output_shape.__doc__
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__

mt/{tf/keras_layers → keras/layers}/image_sizing.py

@@ -1,9 +1,10 @@
 """Module involves upsizing and downsizing images in each axis individually using convolutions of residuals."""
 
-import tensorflow as tf
-
 from mt import tp, np
 
+import tensorflow as tf
+from ..base import layers, initializers, regularizers, constraints
+
 
 def mirror_all_weights(l_weights: list) -> list:
     """TBC"""
@@ -27,7 +28,7 @@ def mirror_all_weights(l_weights: list) -> list:
     return l_newWeights
 
 
-class DUCLayer(tf.keras.layers.Layer):
+class DUCLayer(layers.Layer):
     """Base layer for all DUC layer implementations.
 
     Parameters
@@ -65,33 +66,27 @@ class DUCLayer(tf.keras.layers.Layer):
         super(DUCLayer, self).__init__(**kwargs)
 
         self._kernel_size = kernel_size
-        self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
-        self._bias_initializer = tf.keras.initializers.get(bias_initializer)
-        self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
-        self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
-        self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
-        self._bias_constraint = tf.keras.constraints.get(bias_constraint)
+        self._kernel_initializer = initializers.get(kernel_initializer)
+        self._bias_initializer = initializers.get(bias_initializer)
+        self._kernel_regularizer = regularizers.get(kernel_regularizer)
+        self._bias_regularizer = regularizers.get(bias_regularizer)
+        self._kernel_constraint = constraints.get(kernel_constraint)
+        self._bias_constraint = constraints.get(bias_constraint)
 
     def get_config(self):
         config = {
             "kernel_size": self._kernel_size,
-            "kernel_initializer": tf.keras.initializers.serialize(
-                self._kernel_initializer
-            ),
-            "bias_initializer": tf.keras.initializers.serialize(self._bias_initializer),
-            "kernel_regularizer": tf.keras.regularizers.serialize(
-                self._kernel_regularizer
-            ),
-            "bias_regularizer": tf.keras.regularizers.serialize(self._bias_regularizer),
-            "kernel_constraint": tf.keras.constraints.serialize(
-                self._kernel_constraint
-            ),
-            "bias_constraint": tf.keras.constraints.serialize(self._bias_constraint),
+            "kernel_initializer": initializers.serialize(self._kernel_initializer),
+            "bias_initializer": initializers.serialize(self._bias_initializer),
+            "kernel_regularizer": regularizers.serialize(self._kernel_regularizer),
+            "bias_regularizer": regularizers.serialize(self._bias_regularizer),
+            "kernel_constraint": constraints.serialize(self._kernel_constraint),
+            "bias_constraint": constraints.serialize(self._bias_constraint),
         }
         base_config = super(DUCLayer, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
-    get_config.__doc__ = tf.keras.layers.Layer.get_config.__doc__
+    get_config.__doc__ = layers.Layer.get_config.__doc__
 
     def get_mirrored_weights(self):
         return mirror_all_weights(self.get_weights())
@@ -166,8 +161,8 @@ class Upsize2D(DUCLayer):
         self._expansion_factor = expansion_factor
 
         if self._expansion_factor > 1:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expansion_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expansion_layer = layers.Conv2D(
                 self._input_dim * 2 * expansion_factor,
                 self._kernel_size,
                 padding="same",
@@ -180,8 +175,8 @@ class Upsize2D(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="expand",
             )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.projection_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.projection_layer = layers.Conv2D(
             self._input_dim * 2,
             self._kernel_size,
             padding="same",
@@ -332,8 +327,8 @@ class Downsize2D(DUCLayer):
         self._expansion_factor = expansion_factor
 
         if self._expansion_factor > 1:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expansion_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expansion_layer = layers.Conv2D(
                 self._input_dim * 4 * self._expansion_factor,
                 self._kernel_size,
                 padding="same",
@@ -346,8 +341,8 @@ class Downsize2D(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="expand",
             )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.projection_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.projection_layer = layers.Conv2D(
             self._input_dim,
             self._kernel_size,
             padding="same",
@@ -523,8 +518,8 @@ class Downsize2D_V2(DUCLayer):
         self._projection_uses_bias = projection_uses_bias
 
         if self._expansion_factor > 1:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expansion_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expansion_layer = layers.Conv2D(
                 (self._img_dim + self._res_dim) * 4 * self._expansion_factor,
                 self._kernel_size,
                 padding="same",
@@ -537,8 +532,8 @@ class Downsize2D_V2(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="expand",
             )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.projection_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.projection_layer = layers.Conv2D(
             self._img_dim + self._res_dim * 2,
             1,
             padding="same",
@@ -710,8 +705,8 @@ class Upsize2D_V2(DUCLayer):
         self._expansion_factor = expansion_factor
 
         if self._expansion_factor > 1:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expansion_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expansion_layer = layers.Conv2D(
                 (self._img_dim + self._res_dim) * 2 * expansion_factor,
                 self._kernel_size,
                 padding="same",
@@ -724,8 +719,8 @@ class Upsize2D_V2(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="expand",
             )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.projection_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.projection_layer = layers.Conv2D(
             (self._img_dim + self._res_dim) * 2,
             self._kernel_size if self._expansion_factor <= 1 else 1,
             padding="same",
@@ -875,10 +870,8 @@ class Downsize2D_V3(DUCLayer):
 
         if res_dim > 0:
             if res_dim > img_dim:
-                self.prenorm1_layer = tf.keras.layers.LayerNormalization(
-                    name="prenorm1"
-                )
-                self.expand1_layer = tf.keras.layers.Conv2D(
+                self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+                self.expand1_layer = layers.Conv2D(
                     img_dim * 2 + res_dim * 4,
                     self._kernel_size,
                     padding="same",
@@ -892,8 +885,8 @@ class Downsize2D_V3(DUCLayer):
                     name="expand1",
                 )
             RR = (img_dim + res_dim * 3 + 1) // 2
-            self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-            self.project1_layer = tf.keras.layers.Conv2D(
+            self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+            self.project1_layer = layers.Conv2D(
                 RR,
                 1,
                 padding="same",
@@ -906,8 +899,8 @@ class Downsize2D_V3(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="project1",
             )
-            self.prenorm3_layer = tf.keras.layers.LayerNormalization(name="prenorm3")
-            self.expand2_layer = tf.keras.layers.Conv2D(
+            self.prenorm3_layer = layers.LayerNormalization(name="prenorm3")
+            self.expand2_layer = layers.Conv2D(
                 img_dim * 2 + RR * 4,
                 self._kernel_size,
                 padding="same",
@@ -920,8 +913,8 @@ class Downsize2D_V3(DUCLayer):
                 bias_constraint=self._bias_constraint,
                 name="expand2",
             )
-        self.prenorm4_layer = tf.keras.layers.LayerNormalization(name="prenorm4")
-        self.project2_layer = tf.keras.layers.Conv2D(
+        self.prenorm4_layer = layers.LayerNormalization(name="prenorm4")
+        self.project2_layer = layers.Conv2D(
             img_dim + res_dim * 2,
             1,
             padding="same",
@@ -1174,10 +1167,8 @@ class DownsizeX2D(DUCLayerV5):
 
         if res_dim > 0:
             if res_dim > img_dim:
-                self.prenorm1_layer = tf.keras.layers.LayerNormalization(
-                    name="prenorm1"
-                )
-                self.expand1_layer = tf.keras.layers.Conv2D(
+                self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+                self.expand1_layer = layers.Conv2D(
                     (self.I + self.R) * 4,
                     self._kernel_size,
                     padding="same",
@@ -1190,8 +1181,8 @@ class DownsizeX2D(DUCLayerV5):
                     bias_constraint=self._bias_constraint,
                     name="expand1",
                 )
-            self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-            self.project1_layer = tf.keras.layers.Conv2D(
+            self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+            self.project1_layer = layers.Conv2D(
                 self.RX,
                 1,
                 padding="same",
@@ -1298,8 +1289,8 @@ class UpsizeX2D(DUCLayerV5):
         )
 
         if res_dim > 0:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expand1_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expand1_layer = layers.Conv2D(
                 (self.I + self.R) * 4,
                 self._kernel_size,
                 padding="same",
@@ -1312,8 +1303,8 @@ class UpsizeX2D(DUCLayerV5):
                 bias_constraint=self._bias_constraint,
                 name="expand1",
             )
-            self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-            self.project1_layer = tf.keras.layers.Conv2D(
+            self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+            self.project1_layer = layers.Conv2D(
                 self.R,
                 1,
                 padding="same",
@@ -1326,7 +1317,7 @@ class UpsizeX2D(DUCLayerV5):
                 bias_constraint=self._bias_constraint,
                 name="project1",
             )
-            self.project2_layer = tf.keras.layers.Conv2D(
+            self.project2_layer = layers.Conv2D(
                 self.I + self.R,
                 1,
                 padding="same",
@@ -1429,8 +1420,8 @@ class DownsizeY2D(DUCLayerV5):
         )
 
         if self.R > 0:
-            self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-            self.expand1_layer = tf.keras.layers.Conv2D(
+            self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+            self.expand1_layer = layers.Conv2D(
                 (self.I + self.RX) * 4,
                 self._kernel_size,
                 padding="same",
@@ -1443,8 +1434,8 @@ class DownsizeY2D(DUCLayerV5):
                 bias_constraint=self._bias_constraint,
                 name="expand1",
             )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.project1_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.project1_layer = layers.Conv2D(
             self.RY,
             1,
             padding="same",
@@ -1547,8 +1538,8 @@ class UpsizeY2D(DUCLayerV5):
             **kwargs
         )
 
-        self.prenorm1_layer = tf.keras.layers.LayerNormalization(name="prenorm1")
-        self.expand1_layer = tf.keras.layers.Conv2D(
+        self.prenorm1_layer = layers.LayerNormalization(name="prenorm1")
+        self.expand1_layer = layers.Conv2D(
             (self.I + self.RX) * 4,
             self._kernel_size,
             padding="same",
@@ -1561,8 +1552,8 @@ class UpsizeY2D(DUCLayerV5):
             bias_constraint=self._bias_constraint,
             name="expand1",
         )
-        self.prenorm2_layer = tf.keras.layers.LayerNormalization(name="prenorm2")
-        self.project1_layer = tf.keras.layers.Conv2D(
+        self.prenorm2_layer = layers.LayerNormalization(name="prenorm2")
+        self.project1_layer = layers.Conv2D(
             self.RX,
             1,
             padding="same",
@@ -1575,7 +1566,7 @@ class UpsizeY2D(DUCLayerV5):
             bias_constraint=self._bias_constraint,
             name="project1",
         )
-        self.project2_layer = tf.keras.layers.Conv2D(
+        self.project2_layer = layers.Conv2D(
             self.I + self.RX,
             1,
             padding="same",

mt/{tf/keras_layers → keras/layers}/normed_conv2d.py

@@ -1,10 +1,8 @@
-import tensorflow.keras as tk
-
-
+from ..base import layers, initializers, regularizers, constraints
 from .counter import Counter
 
 
-class NormedConv2D(tk.layers.Layer):
+class NormedConv2D(layers.Layer):
     """A block of Conv2D without activation, followed by LayerNormalization, then activation.
 
     This layer represents the following block:
@@ -171,7 +169,7 @@ class NormedConv2D(tk.layers.Layer):
         for key in self.keys:
             setattr(self, key, locals()[key])
 
-        self.conv2d = tk.layers.Conv2D(
+        self.conv2d = layers.Conv2D(
             filters,
             kernel_size,
             strides=strides,
@@ -188,7 +186,7 @@ class NormedConv2D(tk.layers.Layer):
 
         self.counter = Counter()
 
-        self.norm = tk.layers.LayerNormalization(
+        self.norm = layers.LayerNormalization(
             axis=-1,
             epsilon=epsilon,
             scale=True,
@@ -202,7 +200,7 @@ class NormedConv2D(tk.layers.Layer):
         )
 
         if activation is not None:
-            self.acti = tk.layers.Activation(activation)
+            self.acti = layers.Activation(activation)
 
     def call(self, x, training: bool = False):
         count = self.counter(x, training=training)
@@ -215,7 +213,7 @@ class NormedConv2D(tk.layers.Layer):
         w = self.acti(z, training=training)
         return w
 
-    call.__doc__ = tk.layers.Layer.call.__doc__
+    call.__doc__ = layers.Layer.call.__doc__
 
     def get_config(self):
         config = {key: getattr(self, key) for key in self.keys}
@@ -224,17 +222,17 @@ class NormedConv2D(tk.layers.Layer):
             key = prefix + "_initializer"
             value = config[key]
             if not isinstance(value, str):
-                value = tk.initializers.serialize(value)
+                value = initializers.serialize(value)
             config[key] = value
             key = prefix + "_regularizer"
             value = config[key]
             if not isinstance(value, str):
-                value = tk.regularizers.serialize(value)
+                value = regularizers.serialize(value)
             config[key] = value
             key = prefix + "_constraint"
             value = config[key]
             if not isinstance(value, str):
-                value = tk.constraints.serialize(value)
+                value = constraints.serialize(value)
             config[key] = value
         config = {key: value for key, value in config.items() if value is not None}
         base_config = super(NormedConv2D, self).get_config()

mt/{tf/keras_layers → keras/layers}/simple_mha.py

@@ -21,9 +21,11 @@ from tensorflow.python.util.tf_export import keras_export
 
 from mt import tp, tfc
 
+from ..base import layers, initializers, regularizers, constraints
+
 
 @keras_export("keras.layers.SimpleMHA2D")
-class SimpleMHA2D(tf.keras.layers.Layer):
+class SimpleMHA2D(layers.Layer):
     """SimpleMHA2D layer.
 
     This is a simplified version of the Keras-based MultiHeadAttention layer.
@@ -70,7 +72,7 @@ class SimpleMHA2D(tf.keras.layers.Layer):
     --------
 
     >>> layer = SimpleMHA2D(num_heads=3, key_dim=40, value_dim=80)
-    >>> input_tensor = tf.keras.Input(shape=[8, 8, 160])
+    >>> input_tensor = layers.Input(shape=[8, 8, 160])
     >>> output_tensor = layer(input_tensor)
     >>> print(output_tensor.shape)
     (None, 3, 80)
@@ -97,13 +99,13 @@ class SimpleMHA2D(tf.keras.layers.Layer):
         self._key_dim = key_dim
         self._value_dim = value_dim if value_dim else key_dim
         self._use_bias = use_bias
-        self._activation = tf.keras.activations.get(activation)
-        self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
-        self._bias_initializer = tf.keras.initializers.get(bias_initializer)
-        self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
-        self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
-        self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
-        self._bias_constraint = tf.keras.constraints.get(bias_constraint)
+        self._activation = activations.get(activation)
+        self._kernel_initializer = initializers.get(kernel_initializer)
+        self._bias_initializer = initializers.get(bias_initializer)
+        self._kernel_regularizer = regularizers.get(kernel_regularizer)
+        self._bias_regularizer = regularizers.get(bias_regularizer)
+        self._kernel_constraint = constraints.get(kernel_constraint)
+        self._bias_constraint = constraints.get(bias_constraint)
         self._dropout = dropout
 
         self.tensor_query = self.add_weight(
@@ -113,7 +115,7 @@ class SimpleMHA2D(tf.keras.layers.Layer):
             trainable=True,
         )
 
-        self.layer_key_proj = tf.keras.layers.Conv2D(
+        self.layer_key_proj = layers.Conv2D(
             self._num_heads * self._key_dim,  # filters
             1,  # kernel_size
             use_bias=self._use_bias,
@@ -125,7 +127,7 @@ class SimpleMHA2D(tf.keras.layers.Layer):
             bias_constraint=self._bias_constraint,
         )
 
-        self.layer_value_proj = tf.keras.layers.Conv2D(
+        self.layer_value_proj = layers.Conv2D(
             self._num_heads * self._value_dim,  # filters
             1,  # kernel_size
             use_bias=self._use_bias,
@@ -138,9 +140,9 @@ class SimpleMHA2D(tf.keras.layers.Layer):
             bias_constraint=self._bias_constraint,
         )
 
-        self.layer_softmax = tf.keras.layers.Softmax(axis=1)
+        self.layer_softmax = layers.Softmax(axis=1)
         if self._dropout > 0:
-            self.layer_dropout = tf.keras.layers.Dropout(rate=self._dropout)
+            self.layer_dropout = layers.Dropout(rate=self._dropout)
 
     def call(self, key_value, training=None):
         """The call function.
@@ -203,19 +205,13 @@ class SimpleMHA2D(tf.keras.layers.Layer):
             "key_dim": self._key_dim,
             "value_dim": self._value_dim,
             "use_bias": self._use_bias,
-            "activation": tf.keras.activations.serialize(self._activation),
-            "kernel_initializer": tf.keras.initializers.serialize(
-                self._kernel_initializer
-            ),
-            "bias_initializer": tf.keras.initializers.serialize(self._bias_initializer),
-            "kernel_regularizer": tf.keras.regularizers.serialize(
-                self._kernel_regularizer
-            ),
-            "bias_regularizer": tf.keras.regularizers.serialize(self._bias_regularizer),
-            "kernel_constraint": tf.keras.constraints.serialize(
-                self._kernel_constraint
-            ),
-            "bias_constraint": tf.keras.constraints.serialize(self._bias_constraint),
+            "activation": activations.serialize(self._activation),
+            "kernel_initializer": initializers.serialize(self._kernel_initializer),
+            "bias_initializer": initializers.serialize(self._bias_initializer),
+            "kernel_regularizer": regularizers.serialize(self._kernel_regularizer),
+            "bias_regularizer": regularizers.serialize(self._bias_regularizer),
+            "kernel_constraint": constraints.serialize(self._kernel_constraint),
+            "bias_constraint": constraints.serialize(self._bias_constraint),
             "dropout": self._dropout,
         }
         base_config = super(SimpleMHA2D, self).get_config()
@@ -223,7 +219,7 @@ class SimpleMHA2D(tf.keras.layers.Layer):
 
 
 @keras_export("keras.layers.MHAPool2D")
-class MHAPool2D(tf.keras.layers.Layer):
+class MHAPool2D(layers.Layer):
     """Pooling in 2D using Keras-based self-attention.
 
     The layer takes as input a high-dim image tensor of shape [B, H, W, D] where B is the
@@ -279,7 +275,7 @@ class MHAPool2D(tf.keras.layers.Layer):
     --------
 
     >>> layer = MHAPool2D(num_heads=3, key_dim=40, value_dim=80)
-    >>> input_tensor = tf.keras.Input(shape=[8, 8, 160])
+    >>> input_tensor = layers.Input(shape=[8, 8, 160])
     >>> output_tensor = layer(input_tensor)
     >>> print(output_tensor.shape)
     (None, 4, 4, 240)
@@ -310,25 +306,25 @@ class MHAPool2D(tf.keras.layers.Layer):
         self._pooling = pooling
         self._pool_size = pool_size
         self._use_bias = use_bias
-        self._activation = tf.keras.activations.get(activation)
-        self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
-        self._bias_initializer = tf.keras.initializers.get(bias_initializer)
-        self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
-        self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
-        self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
-        self._bias_constraint = tf.keras.constraints.get(bias_constraint)
+        self._activation = activations.get(activation)
+        self._kernel_initializer = initializers.get(kernel_initializer)
+        self._bias_initializer = initializers.get(bias_initializer)
+        self._kernel_regularizer = regularizers.get(kernel_regularizer)
+        self._bias_regularizer = regularizers.get(bias_regularizer)
+        self._kernel_constraint = constraints.get(kernel_constraint)
+        self._bias_constraint = constraints.get(bias_constraint)
         self._dropout = dropout
 
         if self._pooling == "max":
-            self.layer_pool = tf.keras.layers.MaxPool2D()
+            self.layer_pool = layers.MaxPool2D()
         elif self._pooling == "avg":
-            self.layer_pool = tf.keras.layers.AveragePooling2D()
+            self.layer_pool = layers.AveragePooling2D()
         else:
             raise tfc.ModelSyntaxError(
                 "Invalid pooling string: '{}'.".format(self._pooling)
             )
 
-        self.layer_query_proj = tf.keras.layers.Conv2D(
+        self.layer_query_proj = layers.Conv2D(
             self._num_heads * self._key_dim,  # filters
             1,  # kernel_size
             use_bias=self._use_bias,
@@ -342,7 +338,7 @@ class MHAPool2D(tf.keras.layers.Layer):
             name="query_proj",
         )
 
-        self.layer_key_proj = tf.keras.layers.Conv2D(
+        self.layer_key_proj = layers.Conv2D(
             self._num_heads * self._key_dim,  # filters
             1,  # kernel_size
             use_bias=self._use_bias,
@@ -356,7 +352,7 @@ class MHAPool2D(tf.keras.layers.Layer):
             name="key_proj",
         )
 
-        self.layer_value_proj = tf.keras.layers.Conv2D(
+        self.layer_value_proj = layers.Conv2D(
             self._num_heads * self._value_dim,  # filters
             1,  # kernel_size
             use_bias=self._use_bias,
@@ -370,9 +366,9 @@ class MHAPool2D(tf.keras.layers.Layer):
             name="value_proj",
         )
 
-        self.layer_softmax = tf.keras.layers.Softmax(axis=3)
+        self.layer_softmax = layers.Softmax(axis=3)
         if self._dropout > 0:
-            self.layer_dropout = tf.keras.layers.Dropout(rate=self._dropout)
+            self.layer_dropout = layers.Dropout(rate=self._dropout)
 
     def call(self, blob, training=None, return_attention_scores: bool = False):
         """The call function.
@@ -464,19 +460,13 @@ class MHAPool2D(tf.keras.layers.Layer):
             "pooling": self._pooling,
             "pool_size": self._pool_size,
             "use_bias": self._use_bias,
-            "activation": tf.keras.activations.serialize(self._activation),
-            "kernel_initializer": tf.keras.initializers.serialize(
-                self._kernel_initializer
-            ),
-            "bias_initializer": tf.keras.initializers.serialize(self._bias_initializer),
-            "kernel_regularizer": tf.keras.regularizers.serialize(
-                self._kernel_regularizer
-            ),
-            "bias_regularizer": tf.keras.regularizers.serialize(self._bias_regularizer),
-            "kernel_constraint": tf.keras.constraints.serialize(
-                self._kernel_constraint
-            ),
-            "bias_constraint": tf.keras.constraints.serialize(self._bias_constraint),
+            "activation": activations.serialize(self._activation),
+            "kernel_initializer": initializers.serialize(self._kernel_initializer),
+            "bias_initializer": initializers.serialize(self._bias_initializer),
+            "kernel_regularizer": regularizers.serialize(self._kernel_regularizer),
+            "bias_regularizer": regularizers.serialize(self._bias_regularizer),
+            "kernel_constraint": constraints.serialize(self._kernel_constraint),
+            "bias_constraint": constraints.serialize(self._bias_constraint),
             "dropout": self._dropout,
         }
         base_config = super(MHAPool2D, self).get_config()

mt/{tf/keras_layers → keras/layers}/utils.py

@@ -2,7 +2,7 @@
 
 from mt import tp
 
-from ..utils import NameScope
+from mt.tf.utils import NameScope
 
 
 def conv2d(name_scope: NameScope, x, filters, kernel_size, **kwargs):
@@ -30,11 +30,11 @@ def conv2d(name_scope: NameScope, x, filters, kernel_size, **kwargs):
         TF tensor as output
     """
 
-    import tensorflow.keras.layers as kl
+    from ..base import layers
 
     next(name_scope)
-    x = kl.LayerNormalization(name=name_scope("prenorm"))(x)
-    x = kl.Conv2D(filters, kernel_size, name=name_scope("conv"), **kwargs)(x)
+    x = layers.LayerNormalization(name=name_scope("prenorm"))(x)
+    x = layers.Conv2D(filters, kernel_size, name=name_scope("conv"), **kwargs)(x)
 
     return x
 
@@ -71,15 +71,15 @@ def dense2d(
         TF tensor as output
     """
 
-    import tensorflow.keras.layers as kl
+    from ..base import layers
 
     next(name_scope)
-    x = kl.LayerNormalization(name=name_scope("expand_prenorm"))(x)
-    x = kl.Conv2D(
+    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 = kl.LayerNormalization(name=name_scope("project_prenorm"))(x)
-    x = kl.Conv2D(
+    x = layers.LayerNormalization(name=name_scope("project_prenorm"))(x)
+    x = layers.Conv2D(
         filters,
         kernel_size,
         name=name_scope("project"),

mt/{tf/keras_layers → keras/layers}/var_regularizer.py

@@ -2,8 +2,10 @@ import tensorflow as tf
 
 from mt import tp
 
+from ..base import layers
 
-class VarianceRegularizer(tf.keras.layers.Layer):
+
+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.
@@ -23,7 +25,7 @@ class VarianceRegularizer(tf.keras.layers.Layer):
         self.add_loss(self.rate * sum_var)
         return x
 
-    call.__doc__ = tf.keras.layers.Layer.call.__doc__
+    call.__doc__ = layers.Layer.call.__doc__
 
     def get_config(self):
         config = {
@@ -33,4 +35,4 @@ class VarianceRegularizer(tf.keras.layers.Layer):
         base_config = super(VarianceRegularizer, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
-    get_config.__doc__ = tf.keras.layers.Layer.get_config.__doc__
+    get_config.__doc__ = layers.Layer.get_config.__doc__

mt/tf/keras_layers/__init__.py

@@ -1,31 +1,5 @@
-from .identical import *
-from .floor import *
-from .var_regularizer import *
-from .simple_mha import *
-from .image_sizing import *
-from .counter import Counter
-from .normed_conv2d import NormedConv2D
-from .utils import *
+from mt.logg import logger
 
+logger.warn_module_move("mt.tf.keras_layers", "mt.keras.layers")
 
-__api__ = [
-    "Identical",
-    "Floor",
-    "VarianceRegularizer",
-    "SimpleMHA2D",
-    "MHAPool2D",
-    "DUCLayer",
-    "Downsize2D",
-    "Upsize2D",
-    "Downsize2D_V2",
-    "Upsize2D_V2",
-    "Downsize2D_V3",
-    "Downsize2D_V4",
-    "DownsizeX2D",
-    "UpsizeX2D",
-    "DownsizeY2D",
-    "UpsizeY2D",
-    "Counter",
-    "conv2d",
-    "dense2d",
-]
+from mt.keras.layers import *

mt/tf/version.py

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

{mttf-1.1.5.dist-info → mttf-1.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mttf
-Version: 1.1.5
+Version: 1.1.9
 Summary: A package to detect and monkey-patch TensorFlow and Keras, for Minh-Tri Pham
 Home-page: https://github.com/inteplus/mttf
 Author: ['Minh-Tri Pham']

mttf-1.1.9.dist-info/RECORD

@@ -0,0 +1,36 @@
+mt/keras/__init__.py,sha256=oM2xsSZTuGSUvFJSpntk8r7bzGcTytNvn-Apv_SBFh4,70
+mt/keras/base/__init__.py,sha256=vIvuf3gjbpXC-rGpXLfIApLRUV6w1GUzVK8YJOgWLyk,327
+mt/keras/base/base.py,sha256=_B2sSUMlHOtGSAqQD1p5YD0raEDL4W0Bh3uKD6BXOJM,807
+mt/keras/layers/__init__.py,sha256=mGO1iwfYK9ml1gcvpBZC-fpB_Kc4ayeKKc_qWX3kZGI,757
+mt/keras/layers/counter.py,sha256=J3__IXbaa7zp72a5P8FFi0bfftTHwa1xzzCwxCIU2gc,856
+mt/keras/layers/floor.py,sha256=4mSpmTrhM7VqTK85POkjC3OhaTZUNUF9knO7gTbSGtc,512
+mt/keras/layers/identical.py,sha256=AIqC36PxU9sXyF9rZuQ-5ObjnIjBiSIMHIb5MwqVdmY,361
+mt/keras/layers/image_sizing.py,sha256=LeWwyFwQSYrDq_2fjOVi1YiRn7GcwMzrLUCLR8zcaKI,58530
+mt/keras/layers/normed_conv2d.py,sha256=dFqeuNS0WyzrGIP3wiaKl_iSFmTgTPg-w3nCvn2X4d0,10680
+mt/keras/layers/simple_mha.py,sha256=ecGL5Sy8mbpj6F8hAAN2h-7Dt0TiAOEM7o6SyVrVTp4,19225
+mt/keras/layers/utils.py,sha256=lk9y0Sl4_w69JtFXKrKlWIgWBJx5lz9WrQi2LszvaZE,2834
+mt/keras/layers/var_regularizer.py,sha256=EZ8ueXrObfu-oo4qixM0UkULe-C03kdNQppQ_NmtYSA,1108
+mt/tf/__init__.py,sha256=M8xiJNdrAUJZgiZTOQOdfkehjO-CYzGpoxh5HVGBkms,338
+mt/tf/init.py,sha256=BUpw33uyA_DmeJjrN2fX4MIs8MynKxkwgc2oTGTqssU,1294
+mt/tf/mttf_version.py,sha256=ha53i-H9pE-crufFttUECgXHwPvam07zMKzApUts1Gs,206
+mt/tf/utils.py,sha256=Copl5VM0PpuFUchK-AcBuGO6QitDwHcEs4FruZb2GAI,2460
+mt/tf/version.py,sha256=OJ5C8dPfTjHQppzvBtr3tex8vKiblUMZ8i5wkIS8_U0,206
+mt/tf/keras_applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mt/tf/keras_applications/mobilenet_v3_split.py,sha256=1oPB3EX3k3c7iju9Ksuw9xyv32-mOKPs3uy2Mk5tLd8,19716
+mt/tf/keras_applications/mobilevit.py,sha256=VsKB_U-f9jwUEjpd0eq-YXL4rDsuAbKQ0yIzkbMfLzw,9949
+mt/tf/keras_layers/__init__.py,sha256=NsuFD-kSuy6cVV3Kl7ab95tw4g7x4Igv3cF-Ky3VuCo,124
+mt/tfc/__init__.py,sha256=XFXAbmEi5uTAuZj797cBqw9ZBnVy6ptxP0TTzVauMNk,8446
+mt/tfg/__init__.py,sha256=6Ly2QImAyQTsg_ZszuAuK_L2n56v89Cix9yYmMVk0CM,304
+mt/tfp/__init__.py,sha256=AQkGCkmDRwswEt3qoOSpxe-fZekx78sHHBs2ZVz33gc,383
+mt/tfp/real_nvp.py,sha256=U9EmkXGqFcvtS2yeh5_RgbKlVKKlGFGklAb7Voyazz4,4440
+mttf-1.1.9.data/scripts/dmt_pipi.sh,sha256=NNsj4P332unHMqU4mAFjU9PQvxQ8TK5XQ42LC29IZY8,510
+mttf-1.1.9.data/scripts/dmt_twineu.sh,sha256=KZhcYwuCW0c36tWcOgCe7uxJmS08rz-J6YNY76Exy4M,193
+mttf-1.1.9.data/scripts/pipi.sh,sha256=kdo96bdaKq2QIa52Z4XFSiGPcbDm09SAU9cju6I2Lxo,289
+mttf-1.1.9.data/scripts/wml_nexus.py,sha256=kW0ju8_kdXc4jOjhdzKiMsFuO1MNpHmu87skrhu9SEg,1492
+mttf-1.1.9.data/scripts/wml_pipi.sh,sha256=CuidIcbuxyXSBNQqYRhCcSC8QbBaSGnQX0KAIFaIvKA,499
+mttf-1.1.9.data/scripts/wml_twineu.sh,sha256=av1JLN765oOWC5LPkv2eSWIVof26y60601tMGkuYdb8,180
+mttf-1.1.9.dist-info/licenses/LICENSE,sha256=e_JtcszdGZ2ZGfjcymTGrcxFj_9XPicZOVtnsrPvruk,1070
+mttf-1.1.9.dist-info/METADATA,sha256=YyP4zCEwqSJvYI-6QfagU_NVwuVjHhugu0Jpviq22t0,567
+mttf-1.1.9.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
+mttf-1.1.9.dist-info/top_level.txt,sha256=WcqGFu9cV7iMZg09iam8eNxUvGpLSKKF2Iubf6SJVOo,3
+mttf-1.1.9.dist-info/RECORD,,

mttf-1.1.5.dist-info/RECORD

@@ -1,32 +0,0 @@
-mt/tf/__init__.py,sha256=M8xiJNdrAUJZgiZTOQOdfkehjO-CYzGpoxh5HVGBkms,338
-mt/tf/init.py,sha256=BUpw33uyA_DmeJjrN2fX4MIs8MynKxkwgc2oTGTqssU,1294
-mt/tf/mttf_version.py,sha256=ha53i-H9pE-crufFttUECgXHwPvam07zMKzApUts1Gs,206
-mt/tf/utils.py,sha256=Copl5VM0PpuFUchK-AcBuGO6QitDwHcEs4FruZb2GAI,2460
-mt/tf/version.py,sha256=2_3_n3SpC_2FgUHWv6ZkAnhLtvlhhn8MQuXgUGe03qU,206
-mt/tf/keras_applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mt/tf/keras_applications/mobilenet_v3_split.py,sha256=1oPB3EX3k3c7iju9Ksuw9xyv32-mOKPs3uy2Mk5tLd8,19716
-mt/tf/keras_applications/mobilevit.py,sha256=VsKB_U-f9jwUEjpd0eq-YXL4rDsuAbKQ0yIzkbMfLzw,9949
-mt/tf/keras_layers/__init__.py,sha256=fSfhKmDz4mIHUYXgRrditWY_aAkgWGM_KjmAilOauXg,578
-mt/tf/keras_layers/counter.py,sha256=KsM25lRvxQZZUz0GLbeUNoPLwVoLF1MK0tQPn9-dVWQ,858
-mt/tf/keras_layers/floor.py,sha256=cSP-2B3o50ffTFMHZy-C-r7jWA6OFLMyaxTz9kAlsiM,513
-mt/tf/keras_layers/identical.py,sha256=ZBiK5P0QisWBowVKUODH-LfnJT_YI7MLvhEobIlEGx0,364
-mt/tf/keras_layers/image_sizing.py,sha256=OTfCHTyMSJO3LLOEeCbvJiuyV6FeSvMb5mtvCSYvk9Q,59133
-mt/tf/keras_layers/normed_conv2d.py,sha256=UVV5x6jpjGeOqx6DlSbvdV1O7nO_64-wKT3rXkNQxGA,10669
-mt/tf/keras_layers/simple_mha.py,sha256=Xv4lyFxDWRKYLmsPUUNJeZQYi8rLya7HlJsaXGbmV0E,19710
-mt/tf/keras_layers/utils.py,sha256=PzlrGrtgPMzTaNuY3p3QwnVSqsMnZ6NIbRSEqXEamGA,2828
-mt/tf/keras_layers/var_regularizer.py,sha256=Afe5Mpd8TynsXVD06RtEfFatuSKmruKcNze4_C2E2po,1108
-mt/tfc/__init__.py,sha256=XFXAbmEi5uTAuZj797cBqw9ZBnVy6ptxP0TTzVauMNk,8446
-mt/tfg/__init__.py,sha256=6Ly2QImAyQTsg_ZszuAuK_L2n56v89Cix9yYmMVk0CM,304
-mt/tfp/__init__.py,sha256=AQkGCkmDRwswEt3qoOSpxe-fZekx78sHHBs2ZVz33gc,383
-mt/tfp/real_nvp.py,sha256=U9EmkXGqFcvtS2yeh5_RgbKlVKKlGFGklAb7Voyazz4,4440
-mttf-1.1.5.data/scripts/dmt_pipi.sh,sha256=NNsj4P332unHMqU4mAFjU9PQvxQ8TK5XQ42LC29IZY8,510
-mttf-1.1.5.data/scripts/dmt_twineu.sh,sha256=KZhcYwuCW0c36tWcOgCe7uxJmS08rz-J6YNY76Exy4M,193
-mttf-1.1.5.data/scripts/pipi.sh,sha256=kdo96bdaKq2QIa52Z4XFSiGPcbDm09SAU9cju6I2Lxo,289
-mttf-1.1.5.data/scripts/wml_nexus.py,sha256=kW0ju8_kdXc4jOjhdzKiMsFuO1MNpHmu87skrhu9SEg,1492
-mttf-1.1.5.data/scripts/wml_pipi.sh,sha256=CuidIcbuxyXSBNQqYRhCcSC8QbBaSGnQX0KAIFaIvKA,499
-mttf-1.1.5.data/scripts/wml_twineu.sh,sha256=av1JLN765oOWC5LPkv2eSWIVof26y60601tMGkuYdb8,180
-mttf-1.1.5.dist-info/licenses/LICENSE,sha256=e_JtcszdGZ2ZGfjcymTGrcxFj_9XPicZOVtnsrPvruk,1070
-mttf-1.1.5.dist-info/METADATA,sha256=x3hD-T6kgdeBoB6EKTrWmrTPrmL_uM9NOs5q_Hed_Ns,567
-mttf-1.1.5.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-mttf-1.1.5.dist-info/top_level.txt,sha256=WcqGFu9cV7iMZg09iam8eNxUvGpLSKKF2Iubf6SJVOo,3
-mttf-1.1.5.dist-info/RECORD,,