From 35e3b72d5d92e363735cdfb33071cb58508d66e1 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 28 Jan 2024 19:15:24 +0900 Subject: [PATCH] fix: add MultiHeadAttention's __call__ to the stubtest allowlist --- stubs/tensorflow/@tests/stubtest_allowlist.txt | 2 ++ stubs/tensorflow/tensorflow/keras/layers/__init__.pyi | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 5700e283375f..b9fb018ede9d 100644 --- 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 index 7d6d20b17dae..a683ccb1cf00 100644 --- 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