remove MultiHeadAttention...

python · JelleZijlstra · Mar 13, 2024 · Jan 28, 2024 · Jan 28, 2024 · Jan 28, 2024
commit 6e83aa66f2797cd2359aadb9d5e2a55848bedd16
diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -60,7 +60,6 @@ tensorflow.keras.layers.*.__init__
 tensorflow.keras.layers.*.call
 tensorflow.keras.regularizers.Regularizer.__call__
 tensorflow.keras.constraints.Constraint.__call__
-tensorflow.keras.layers.MultiHeadAttention.__call__
 
 # Layer class does good deal of __new__ magic and actually returns one of two different internal
 # types depending on tensorflow execution mode. This feels like implementation internal.

diff --git a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi
@@ -10,7 +10,7 @@ from tensorflow.keras.activations import _Activation
 from tensorflow.keras.constraints import Constraint
 from tensorflow.keras.initializers import _Initializer
 from tensorflow.keras.layers.preprocessing import IntegerLookup as IntegerLookup, StringLookup as StringLookup
-from tensorflow.keras.regularizers import Regularizer, _Regularizer
+from tensorflow.keras.regularizers import _Regularizer
 from tensorflow.python.feature_column.feature_column_v2 import DenseColumn, SequenceDenseColumn
 
 _InputT = TypeVar("_InputT", contravariant=True)
@@ -255,62 +255,6 @@ class DenseFeatures(Layer[Mapping[str, TensorLike], tf.Tensor]):
         name: str | None = None,
     ) -> None: ...
 
-class MultiHeadAttention(Layer[Any, tf.Tensor]):
-    def __init__(
-        self,
-        num_heads: int,
-        key_dim: int | None,
-        value_dim: int | None = None,
-        dropout: float = 0.0,
-        use_bias: bool = True,
-        output_shape: tuple[int, ...] | None = None,
-        attention_axes: tuple[int, ...] | None = None,
-        kernel_initialize: _Initializer = "glorot_uniform",
-        bias_initializer: _Initializer = "zeros",
-        kernel_regularizer: Regularizer | None = None,
-        bias_regularizer: _Regularizer | None = None,
-        activity_regularizer: _Regularizer | None = None,
-        kernel_constraint: _Constraint | None = None,
-        bias_constraint: _Constraint | None = None,
-        trainable: bool = True,
-        dtype: _LayerDtype | None = None,
-        dynamic: bool = False,
-        name: str | None = None,
-    ) -> None: ...
-    @overload
-    def __call__(
-        self,
-        query: tf.Tensor,
-        value: tf.Tensor,
-        key: tf.Tensor | None,
-        attention_mask: tf.Tensor | None,
-        return_attention_scores: Literal[False],
-        training: bool,
-        use_causal_mask: bool,
-    ) -> tf.Tensor: ...
-    @overload
-    def __call__(
-        self,
-        query: tf.Tensor,
-        value: tf.Tensor,
-        key: tf.Tensor | None,
-        attention_mask: tf.Tensor | None,
-        return_attention_scores: Literal[True],
-        training: bool,
-        use_causal_mask: bool,
-    ) -> tuple[tf.Tensor, tf.Tensor]: ...
-    @overload
-    def __call__(
-        self,
-        query: tf.Tensor,
-        value: tf.Tensor,
-        key: tf.Tensor | None = None,
-        attention_mask: tf.Tensor | None = None,
-        return_attention_scores: bool = False,
-        training: bool = False,
-        use_causal_mask: bool = False,
-    ) -> tuple[tf.Tensor, tf.Tensor] | tf.Tensor: ...
-
 class GaussianDropout(Layer[tf.Tensor, tf.Tensor]):
     def __init__(
         self,