fix: add MultiHeadAttention's __call__ to the stubtest allowlist

python · Jan 31, 2024 · 35e3b72 · 35e3b72
1 parent 6e26f9b
commit 35e3b72
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -60,6 +60,8 @@ tensorflow.keras.layers.*.__init__
 tensorflow.keras.layers.*.call
 tensorflow.keras.regularizers.Regularizer.__call__
 tensorflow.keras.constraints.Constraint.__call__
+tensorflow.keras.layers.MultiHeadAttention.__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
@@ -5,7 +5,7 @@ from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
 from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization
-from tensorflow._aliases import AnyArray, TensorLike, TensorCompatible, DTypeLike
+from tensorflow._aliases import AnyArray, DTypeLike, TensorCompatible, TensorLike
 from tensorflow.keras.activations import _Activation
 from tensorflow.keras.constraints import Constraint
 from tensorflow.keras.initializers import _Initializer