From ca822c2c25e4e07f6015f6101f0fa88402f796d8 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 28 Jan 2024 18:04:10 +0900 Subject: [PATCH 01/18] Add some missing keras layers --- stubs/tensorflow/tensorflow/keras/layers.pyi | 165 ++++++++++++++++++- 1 file changed, 161 insertions(+), 4 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index 4fa88ede01d8..267fb3d5d97e 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -1,15 +1,16 @@ from _typeshed import Incomplete -from collections.abc import Callable, Iterable, Sequence -from typing import Any, Generic, TypeVar, overload +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import Any, Generic, Literal, TypeVar, overload from typing_extensions import Self, TypeAlias import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization -from tensorflow._aliases import AnyArray, DTypeLike, TensorCompatible +from tensorflow._aliases import AnyArray, TensorLike, TensorCompatible, DTypeLike from tensorflow.keras.activations import _Activation from tensorflow.keras.constraints import Constraint from tensorflow.keras.initializers import _Initializer -from tensorflow.keras.regularizers import _Regularizer +from tensorflow.keras.regularizers import Regularizer, _Regularizer +from tensorflow.python.feature_column.feature_column_v2 import DenseColumn, SequenceDenseColumn _InputT = TypeVar("_InputT", contravariant=True) _OutputT = TypeVar("_OutputT", covariant=True) @@ -194,4 +195,160 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]): name: str | None = None, ) -> None: ... +class Conv2D(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + filters: int, + kernel_size: int | tuple[int, int], + strides: int | tuple[int, int] = (1, 1), + padding: Literal["valid", "same"] = "valid", + data_format: None | Literal["channels_last", "channels_first"] = None, + dilation_rate: int | tuple[int, int] = (1, 1), + groups: int = 1, + activation: _Activation = None, + use_bias: bool = True, + kernel_initializer: _Initializer = "glorot_uniform", + bias_initializer: _Initializer = "zeros", + kernel_regularizer: _Regularizer = None, + bias_regularizer: _Regularizer = None, + activity_regularizer: _Regularizer = None, + kernel_constraint: _Constraint = None, + bias_constraint: _Constraint = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, + name: str | None = None, + ) -> None: ... + +class Identity(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, trainable: bool = True, dtype: _LayerDtype = None, dynamic: bool = False, name: str | None = None + ) -> None: ... + +class LayerNormalization(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + axis: int = -1, + epsilon: float = 0.001, + center: bool = True, + scale: bool = True, + beta_initializer: _Initializer = "zeros", + gamma_initializer: _Initializer = "ones", + beta_regularizer: _Regularizer = None, + gamma_regularizer: _Regularizer = None, + beta_constraint: _Constraint = None, + gamma_constraint: _Constraint = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, + name: str | None = None, + ) -> None: ... + +class _IndexLookup(Layer[TensorLike, TensorLike]): + @overload + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: ... + @overload + def __call__(self, inputs: tf.SparseTensor) -> tf.SparseTensor: ... + @overload + def __call__(self, inputs: tf.RaggedTensor) -> tf.RaggedTensor: ... + def vocabulary_size(self) -> int: ... + +class StringLookup(_IndexLookup): + def __init__( + self, + max_tokens: int | None = None, + num_oov_indices: int = 1, + mask_token: str | None = None, + oov_token: str = "[UNK]", + vocabulary: str | None | TensorCompatible = None, + idf_weights: TensorCompatible | None = None, + encoding: str = "utf-8", + invert: bool = False, + output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", + sparse: bool = False, + pad_to_max_tokens: bool = False, + ) -> None: ... + +class IntegerLookup(_IndexLookup): + def __init__( + self, + max_tokens: int | None = None, + num_oov_indices: int = 1, + mask_token: int | None = None, + oov_token: int = -1, + vocabulary: str | None | TensorCompatible = None, + vocabulary_dtype: Literal["int64", "int32"] = "int64", + idf_weights: TensorCompatible | None = None, + invert: bool = False, + output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", + sparse: bool = False, + pad_to_max_tokens: bool = False, + ) -> None: ... + +class DenseFeatures(Layer[Mapping[str, TensorLike], tf.Tensor]): + def __init__( + self, + feature_columns: Sequence[DenseColumn | SequenceDenseColumn], + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, + 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, + dynamic: bool = False, + name: str | None = None, + ) -> None: ... + @overload + def __call__( + self, + query: tf.Tensor, + value: tf.Tensor, + key: tf.Tensor | None = None, + attention_mask: tf.Tensor | None = None, + return_attention_scores: Literal[False] = False, + training: bool = False, + use_causal_mask: bool = False, + ) -> 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: Literal[True] = True, + training: bool = False, + use_causal_mask: bool = False, + ) -> tuple[tf.Tensor, tf.Tensor]: ... + +class GaussianDropout(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + rate: float, + seed: int | None = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, + name: str | None = None, + ) -> None: ... + def __getattr__(name: str) -> Incomplete: ... From 6e26f9b6b7aaa61c01f302618720c5ec10f0e396 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 28 Jan 2024 18:34:30 +0900 Subject: [PATCH 02/18] Move some layers to keras.layers.preprocessing [pre-commit.ci] auto fixes from pre-commit.com hooks Move some layers to keras.layers.preprocessing fix: fix imports --- .../keras/{layers.pyi => layers/__init__.pyi} | 58 +++---------------- .../layers/experimental/preprocessing.pyi | 6 ++ .../keras/layers/preprocessing/__init__.pyi | 36 ++++++++++++ .../layers/preprocessing/index_lookup.pyi | 9 +++ 4 files changed, 60 insertions(+), 49 deletions(-) rename stubs/tensorflow/tensorflow/keras/{layers.pyi => layers/__init__.pyi} (86%) create mode 100644 stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi similarity index 86% rename from stubs/tensorflow/tensorflow/keras/layers.pyi rename to stubs/tensorflow/tensorflow/keras/layers/__init__.pyi index 267fb3d5d97e..7d6d20b17dae 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi @@ -9,6 +9,7 @@ from tensorflow._aliases import AnyArray, TensorLike, TensorCompatible, DTypeLik 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.python.feature_column.feature_column_v2 import DenseColumn, SequenceDenseColumn @@ -128,7 +129,7 @@ class Dense(Layer[tf.Tensor, tf.Tensor]): kernel_constraint: _Constraint = None, bias_constraint: _Constraint = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -150,7 +151,7 @@ class BatchNormalization(Layer[tf.Tensor, tf.Tensor]): beta_constraint: _Constraint = None, gamma_constraint: _Constraint = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -162,7 +163,7 @@ class ReLU(Layer[tf.Tensor, tf.Tensor]): negative_slope: float | None = 0.0, threshold: float | None = 0.0, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -174,7 +175,7 @@ class Dropout(Layer[tf.Tensor, tf.Tensor]): noise_shape: TensorCompatible | Sequence[int | None] | None = None, seed: int | None = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -190,7 +191,7 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]): mask_zero: bool = False, input_length: int | None = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -215,7 +216,7 @@ class Conv2D(Layer[tf.Tensor, tf.Tensor]): kernel_constraint: _Constraint = None, bias_constraint: _Constraint = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -239,52 +240,11 @@ class LayerNormalization(Layer[tf.Tensor, tf.Tensor]): beta_constraint: _Constraint = None, gamma_constraint: _Constraint = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... -class _IndexLookup(Layer[TensorLike, TensorLike]): - @overload - def __call__(self, inputs: tf.Tensor) -> tf.Tensor: ... - @overload - def __call__(self, inputs: tf.SparseTensor) -> tf.SparseTensor: ... - @overload - def __call__(self, inputs: tf.RaggedTensor) -> tf.RaggedTensor: ... - def vocabulary_size(self) -> int: ... - -class StringLookup(_IndexLookup): - def __init__( - self, - max_tokens: int | None = None, - num_oov_indices: int = 1, - mask_token: str | None = None, - oov_token: str = "[UNK]", - vocabulary: str | None | TensorCompatible = None, - idf_weights: TensorCompatible | None = None, - encoding: str = "utf-8", - invert: bool = False, - output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", - sparse: bool = False, - pad_to_max_tokens: bool = False, - ) -> None: ... - -class IntegerLookup(_IndexLookup): - def __init__( - self, - max_tokens: int | None = None, - num_oov_indices: int = 1, - mask_token: int | None = None, - oov_token: int = -1, - vocabulary: str | None | TensorCompatible = None, - vocabulary_dtype: Literal["int64", "int32"] = "int64", - idf_weights: TensorCompatible | None = None, - invert: bool = False, - output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", - sparse: bool = False, - pad_to_max_tokens: bool = False, - ) -> None: ... - class DenseFeatures(Layer[Mapping[str, TensorLike], tf.Tensor]): def __init__( self, @@ -313,7 +273,7 @@ class MultiHeadAttention(Layer[Any, tf.Tensor]): kernel_constraint: _Constraint | None = None, bias_constraint: _Constraint | None = None, trainable: bool = True, - dtype: _LayerDtype = None, + dtype: _LayerDtype | None = None, dynamic: bool = False, name: str | None = None, ) -> None: ... diff --git a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi new file mode 100644 index 000000000000..482a0aeaab10 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi @@ -0,0 +1,6 @@ +import abc + +from tensorflow._aliases import TensorLike +from tensorflow.keras.layers import Layer + +class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): ... diff --git a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi new file mode 100644 index 000000000000..1896143f5a81 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi @@ -0,0 +1,36 @@ +from typing import Literal + +from tensorflow._aliases import TensorCompatible +from tensorflow.keras.layers.preprocessing.index_lookup import IndexLookup + +class StringLookup(IndexLookup): + def __init__( + self, + max_tokens: int | None = None, + num_oov_indices: int = 1, + mask_token: str | None = None, + oov_token: str = "[UNK]", + vocabulary: str | None | TensorCompatible = None, + idf_weights: TensorCompatible | None = None, + encoding: str = "utf-8", + invert: bool = False, + output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", + sparse: bool = False, + pad_to_max_tokens: bool = False, + ) -> None: ... + +class IntegerLookup(IndexLookup): + def __init__( + self, + max_tokens: int | None = None, + num_oov_indices: int = 1, + mask_token: int | None = None, + oov_token: int = -1, + vocabulary: str | None | TensorCompatible = None, + vocabulary_dtype: Literal["int64", "int32"] = "int64", + idf_weights: TensorCompatible | None = None, + invert: bool = False, + output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int", + sparse: bool = False, + pad_to_max_tokens: bool = False, + ) -> None: ... diff --git a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi new file mode 100644 index 000000000000..3fd543a87e9a --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi @@ -0,0 +1,9 @@ +from _typeshed import Incomplete + +import tensorflow as tf +from tensorflow.keras.layers.experimental.preprocessing import PreprocessingLayer + +class IndexLookup(PreprocessingLayer): + def compute_output_signature(self, input_spec: Incomplete) -> tf.TensorSpec: ... + def get_vocabulary(self, include_special_tokens: bool = True) -> list[Incomplete]: ... + def vocabulary_size(self) -> int: ... From 35056a765304c36cf3231fe6de7be1ac4f8f9d62 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 28 Jan 2024 19:15:24 +0900 Subject: [PATCH 03/18] fix: add MultiHeadAttention's __call__ to the stubtest allowlist attempt to fix MultiHeadAttention's __call__ --- .../tensorflow/@tests/stubtest_allowlist.txt | 1 + .../tensorflow/keras/layers/__init__.pyi | 27 +++++++++++++------ 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 5700e283375f..375e76454a96 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -60,6 +60,7 @@ 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 index 7d6d20b17dae..24e7c0f7bb4c 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 @@ -282,23 +282,34 @@ class MultiHeadAttention(Layer[Any, tf.Tensor]): self, query: tf.Tensor, value: tf.Tensor, - key: tf.Tensor | None = None, - attention_mask: tf.Tensor | None = None, - return_attention_scores: Literal[False] = False, - training: bool = False, - use_causal_mask: bool = False, + 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: Literal[True] = True, + return_attention_scores: bool = False, training: bool = False, use_causal_mask: bool = False, - ) -> tuple[tf.Tensor, tf.Tensor]: ... + ) -> tuple[tf.Tensor, tf.Tensor] | tf.Tensor: ... class GaussianDropout(Layer[tf.Tensor, tf.Tensor]): def __init__( From 6e83aa66f2797cd2359aadb9d5e2a55848bedd16 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Wed, 31 Jan 2024 23:07:45 +0900 Subject: [PATCH 04/18] remove MultiHeadAttention... --- .../tensorflow/@tests/stubtest_allowlist.txt | 1 - .../tensorflow/keras/layers/__init__.pyi | 58 +------------------ 2 files changed, 1 insertion(+), 58 deletions(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 375e76454a96..5700e283375f 100644 --- 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 index 24e7c0f7bb4c..f7d0f5003be9 100644 --- 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, From cd8632a48b8cd583d4a59aeea6ed4d7a60ba033d Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Thu, 1 Feb 2024 12:30:42 +0900 Subject: [PATCH 05/18] Revert "remove MultiHeadAttention..." This reverts commit 6e83aa66f2797cd2359aadb9d5e2a55848bedd16. --- .../tensorflow/@tests/stubtest_allowlist.txt | 1 + .../tensorflow/keras/layers/__init__.pyi | 58 ++++++++++++++++++- 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 5700e283375f..375e76454a96 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -60,6 +60,7 @@ 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 index f7d0f5003be9..24e7c0f7bb4c 100644 --- 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 +from tensorflow.keras.regularizers import Regularizer, _Regularizer from tensorflow.python.feature_column.feature_column_v2 import DenseColumn, SequenceDenseColumn _InputT = TypeVar("_InputT", contravariant=True) @@ -255,6 +255,62 @@ 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, From 755b0ab2f5c700a689893b0b9fcb7043f6a26ddc Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Thu, 1 Feb 2024 12:31:48 +0900 Subject: [PATCH 06/18] type ignore the override --- stubs/tensorflow/tensorflow/keras/layers/__init__.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi index 24e7c0f7bb4c..7929458d35d4 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete from collections.abc import Callable, Iterable, Mapping, Sequence from typing import Any, Generic, Literal, TypeVar, overload -from typing_extensions import Self, TypeAlias +from typing_extensions import Self, TypeAlias, override import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization @@ -277,6 +277,7 @@ class MultiHeadAttention(Layer[Any, tf.Tensor]): dynamic: bool = False, name: str | None = None, ) -> None: ... + @override # type: ignore @overload def __call__( self, From 9eda7f3b9986bfdf72a0c0725df57f899b0606c1 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 4 Feb 2024 17:28:15 +0900 Subject: [PATCH 07/18] try to fix mypy crash. --- stubs/tensorflow/tensorflow/keras/layers/__init__.pyi | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi index 7929458d35d4..fbd7c8e8bca7 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete from collections.abc import Callable, Iterable, Mapping, Sequence from typing import Any, Generic, Literal, TypeVar, overload -from typing_extensions import Self, TypeAlias, override +from typing_extensions import Self, TypeAlias import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization @@ -277,8 +277,8 @@ class MultiHeadAttention(Layer[Any, tf.Tensor]): dynamic: bool = False, name: str | None = None, ) -> None: ... - @override # type: ignore - @overload + # @override + @overload # type: ignore def __call__( self, query: tf.Tensor, From 2c41a03c4f93c7a6dad9f388c51ae1883f932c15 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 5 Feb 2024 22:59:59 +0900 Subject: [PATCH 08/18] add modules to allowlist due to cursed imports. --- stubs/tensorflow/@tests/stubtest_allowlist.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index c46e92aaba27..85c9c1feed3c 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -47,6 +47,9 @@ tensorflow._aliases # but the real module file is completely different name (even package) and dynamically handled. # tf.initializers at runtime is tensorflow.initializers +# Other cursed import magic similar to the one above. +tensorflow.keras.layers.preprocessing +tensorflow.keras.layers.preprocessing.index_lookup # Layer constructor's always have **kwargs, but only allow a few specific values. PEP 692 # would allow us to specify this with **kwargs and remove the need for these exceptions. From 976b9c8e806d32a26274e3eef7cb8481f3de50b0 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Thu, 8 Feb 2024 22:18:14 +0900 Subject: [PATCH 09/18] remove tensorflow.keras.layers.MultiHeadAttention.__call__ from allowlist --- stubs/tensorflow/@tests/stubtest_allowlist.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 85c9c1feed3c..46a2605f36ff 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -69,7 +69,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. From a7739fb1485faf1639afc5a18b0978e857a0ba0b Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 14:06:46 +0900 Subject: [PATCH 10/18] test --- .../feature_column/feature_column_v2.pyi | 21 ++++++++++--------- .../sequence_feature_column.pyi | 13 ++++-------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi b/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi index e2f02fa4b40f..ea47f769345c 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi +++ b/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi @@ -9,6 +9,7 @@ from typing import Literal from typing_extensions import Self, TypeAlias import tensorflow as tf +from tensorflow import DType from tensorflow._aliases import ShapeLike _Combiners: TypeAlias = Literal["mean", "sqrtn", "sum"] @@ -37,7 +38,7 @@ class NumericColumn(DenseColumn): key: str shape: ShapeLike default_value: float - dtype: tf.DType + dtype: DType normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None def __new__( @@ -45,7 +46,7 @@ class NumericColumn(DenseColumn): key: str, shape: ShapeLike, default_value: float, - dtype: tf.DType, + dtype: DType, normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None, ) -> Self: ... @property @@ -174,9 +175,9 @@ class IdentityCategoricalColumn(CategoricalColumn): class HashedCategoricalColumn(CategoricalColumn): key: str hash_bucket_size: int - dtype: tf.DType + dtype: DType - def __new__(_cls, key: str, hash_bucket_size: int, dtype: tf.DType) -> Self: ... + def __new__(_cls, key: str, hash_bucket_size: int, dtype: DType) -> Self: ... @property def name(self) -> str: ... @property @@ -191,7 +192,7 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): vocabulary_file: str vocabulary_size: int | None num_oov_buckets: int - dtype: tf.DType + dtype: DType default_value: str | int | None file_format: str | None @@ -201,7 +202,7 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): vocabulary_file: str, vocabulary_size: int | None, num_oov_buckets: int, - dtype: tf.DType, + dtype: DType, default_value: str | int | None, file_format: str | None = None, ) -> Self: ... @@ -217,12 +218,12 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): class VocabularyListCategoricalColumn(CategoricalColumn): key: str vocabulary_list: Sequence[str] | Sequence[int] - dtype: tf.DType + dtype: DType default_value: str | int | None num_oov_buckets: int def __new__( - _cls, key: str, vocabulary_list: Sequence[str], dtype: tf.DType, default_value: str | int | None, num_oov_buckets: int + _cls, key: str, vocabulary_list: Sequence[str], dtype: DType, default_value: str | int | None, num_oov_buckets: int ) -> Self: ... @property def name(self) -> str: ... @@ -236,9 +237,9 @@ class VocabularyListCategoricalColumn(CategoricalColumn): class WeightedCategoricalColumn(CategoricalColumn): categorical_column: CategoricalColumn weight_feature_key: str - dtype: tf.DType + dtype: DType - def __new__(_cls, categorical_column: CategoricalColumn, weight_feature_key: str, dtype: tf.DType) -> Self: ... + def __new__(_cls, categorical_column: CategoricalColumn, weight_feature_key: str, dtype: DType) -> Self: ... @property def name(self) -> str: ... @property diff --git a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi b/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi index 921ab086518a..5361f452a5e0 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi +++ b/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi @@ -1,7 +1,7 @@ from collections.abc import Callable from typing_extensions import Self -import tensorflow as tf +from tensorflow import DType, Tensor from tensorflow._aliases import ShapeLike from tensorflow.python.feature_column.feature_column_v2 import SequenceDenseColumn, _ExampleSpec, _FeatureColumn @@ -11,16 +11,11 @@ class SequenceNumericColumn(SequenceDenseColumn): key: str shape: ShapeLike default_value: float - dtype: tf.DType - normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None + dtype: DType + normalizer_fn: Callable[[Tensor], Tensor] | None def __new__( - _cls, - key: str, - shape: ShapeLike, - default_value: float, - dtype: tf.DType, - normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None, + _cls, key: str, shape: ShapeLike, default_value: float, dtype: DType, normalizer_fn: Callable[[Tensor], Tensor] | None ) -> Self: ... @property def name(self) -> str: ... From 03c3e9b0f901aa926358db8feb02eb03d62f5949 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 14:59:28 +0900 Subject: [PATCH 11/18] Revert "test" This reverts commit a7739fb1485faf1639afc5a18b0978e857a0ba0b. --- .../feature_column/feature_column_v2.pyi | 21 +++++++++---------- .../sequence_feature_column.pyi | 13 ++++++++---- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi b/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi index ea47f769345c..e2f02fa4b40f 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi +++ b/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi @@ -9,7 +9,6 @@ from typing import Literal from typing_extensions import Self, TypeAlias import tensorflow as tf -from tensorflow import DType from tensorflow._aliases import ShapeLike _Combiners: TypeAlias = Literal["mean", "sqrtn", "sum"] @@ -38,7 +37,7 @@ class NumericColumn(DenseColumn): key: str shape: ShapeLike default_value: float - dtype: DType + dtype: tf.DType normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None def __new__( @@ -46,7 +45,7 @@ class NumericColumn(DenseColumn): key: str, shape: ShapeLike, default_value: float, - dtype: DType, + dtype: tf.DType, normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None, ) -> Self: ... @property @@ -175,9 +174,9 @@ class IdentityCategoricalColumn(CategoricalColumn): class HashedCategoricalColumn(CategoricalColumn): key: str hash_bucket_size: int - dtype: DType + dtype: tf.DType - def __new__(_cls, key: str, hash_bucket_size: int, dtype: DType) -> Self: ... + def __new__(_cls, key: str, hash_bucket_size: int, dtype: tf.DType) -> Self: ... @property def name(self) -> str: ... @property @@ -192,7 +191,7 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): vocabulary_file: str vocabulary_size: int | None num_oov_buckets: int - dtype: DType + dtype: tf.DType default_value: str | int | None file_format: str | None @@ -202,7 +201,7 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): vocabulary_file: str, vocabulary_size: int | None, num_oov_buckets: int, - dtype: DType, + dtype: tf.DType, default_value: str | int | None, file_format: str | None = None, ) -> Self: ... @@ -218,12 +217,12 @@ class VocabularyFileCategoricalColumn(CategoricalColumn): class VocabularyListCategoricalColumn(CategoricalColumn): key: str vocabulary_list: Sequence[str] | Sequence[int] - dtype: DType + dtype: tf.DType default_value: str | int | None num_oov_buckets: int def __new__( - _cls, key: str, vocabulary_list: Sequence[str], dtype: DType, default_value: str | int | None, num_oov_buckets: int + _cls, key: str, vocabulary_list: Sequence[str], dtype: tf.DType, default_value: str | int | None, num_oov_buckets: int ) -> Self: ... @property def name(self) -> str: ... @@ -237,9 +236,9 @@ class VocabularyListCategoricalColumn(CategoricalColumn): class WeightedCategoricalColumn(CategoricalColumn): categorical_column: CategoricalColumn weight_feature_key: str - dtype: DType + dtype: tf.DType - def __new__(_cls, categorical_column: CategoricalColumn, weight_feature_key: str, dtype: DType) -> Self: ... + def __new__(_cls, categorical_column: CategoricalColumn, weight_feature_key: str, dtype: tf.DType) -> Self: ... @property def name(self) -> str: ... @property diff --git a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi b/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi index 5361f452a5e0..921ab086518a 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi +++ b/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi @@ -1,7 +1,7 @@ from collections.abc import Callable from typing_extensions import Self -from tensorflow import DType, Tensor +import tensorflow as tf from tensorflow._aliases import ShapeLike from tensorflow.python.feature_column.feature_column_v2 import SequenceDenseColumn, _ExampleSpec, _FeatureColumn @@ -11,11 +11,16 @@ class SequenceNumericColumn(SequenceDenseColumn): key: str shape: ShapeLike default_value: float - dtype: DType - normalizer_fn: Callable[[Tensor], Tensor] | None + dtype: tf.DType + normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None def __new__( - _cls, key: str, shape: ShapeLike, default_value: float, dtype: DType, normalizer_fn: Callable[[Tensor], Tensor] | None + _cls, + key: str, + shape: ShapeLike, + default_value: float, + dtype: tf.DType, + normalizer_fn: Callable[[tf.Tensor], tf.Tensor] | None, ) -> Self: ... @property def name(self) -> str: ... From 7d0343fc65cc1d051ce4c03301fcf9b426ac47ae Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 19:57:49 +0900 Subject: [PATCH 12/18] fix: tuple -> Iterable --- stubs/tensorflow/tensorflow/keras/layers/__init__.pyi | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi index fbd7c8e8bca7..0cbd3f381afb 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/__init__.pyi @@ -200,11 +200,11 @@ class Conv2D(Layer[tf.Tensor, tf.Tensor]): def __init__( self, filters: int, - kernel_size: int | tuple[int, int], - strides: int | tuple[int, int] = (1, 1), + kernel_size: int | Iterable[int], + strides: int | Iterable[int] = (1, 1), padding: Literal["valid", "same"] = "valid", data_format: None | Literal["channels_last", "channels_first"] = None, - dilation_rate: int | tuple[int, int] = (1, 1), + dilation_rate: int | Iterable[int] = (1, 1), groups: int = 1, activation: _Activation = None, use_bias: bool = True, From 37c668ba1943d116099bedb75d05274a7fdd31b1 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 20:07:13 +0900 Subject: [PATCH 13/18] add PreprocessingLayer methods/overloads --- .../keras/layers/experimental/preprocessing.pyi | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi index 482a0aeaab10..3256201d2887 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi @@ -1,6 +1,18 @@ import abc +from typing import overload -from tensorflow._aliases import TensorLike +import tensorflow as tf +from tensorflow._aliases import AnyArray, TensorLike from tensorflow.keras.layers import Layer -class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): ... +class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): + @overload + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: ... + @overload + def __call__(self, inputs: tf.SparseTensor) -> tf.SparseTensor: ... + @overload + def __call__(self, inputs: tf.RaggedTensor) -> tf.RaggedTensor: ... # type: ignore + def adapt( + self, data: tf.data.Dataset[TensorLike] | AnyArray, batch_size: int | None = None, steps: int | None = None + ) -> None: ... + def compile(self, run_eagerly: bool | None = None, steps_per_execution: int | None = None) -> None: ... From 4228893b9706dca590a71db713a44ec7ebd853ff Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 20:23:51 +0900 Subject: [PATCH 14/18] fix PreprocessingLayer typing --- stubs/tensorflow/tensorflow/_aliases.pyi | 3 +++ .../keras/layers/experimental/preprocessing.pyi | 10 +++++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index 842c5e5acd47..f506e86b112c 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -27,10 +27,13 @@ class KerasSerializable2(Protocol): KerasSerializable: TypeAlias = KerasSerializable1 | KerasSerializable2 +Integer: TypeAlias = tf.Tensor | int | IntArray | np.number[Any] +Float: TypeAlias = Integer | float | FloatArray Slice: TypeAlias = int | slice | None FloatDataSequence: TypeAlias = Sequence[float] | Sequence[FloatDataSequence] IntDataSequence: TypeAlias = Sequence[int] | Sequence[IntDataSequence] StrDataSequence: TypeAlias = Sequence[str] | Sequence[StrDataSequence] +DataSequence: TypeAlias = FloatDataSequence | StrDataSequence | IntDataSequence ScalarTensorCompatible: TypeAlias = tf.Tensor | str | float | np.ndarray[Any, Any] | np.number[Any] UIntTensorCompatible: TypeAlias = tf.Tensor | int | UIntArray StringTensorCompatible: TypeAlias = tf.Tensor | str | npt.NDArray[np.str_] | Sequence[StringTensorCompatible] diff --git a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi index 3256201d2887..431222f18ecd 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi @@ -2,10 +2,11 @@ import abc from typing import overload import tensorflow as tf -from tensorflow._aliases import AnyArray, TensorLike +from tensorflow._aliases import AnyArray, DataSequence, Float, Integer, TensorLike from tensorflow.keras.layers import Layer class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): + is_adapted: bool @overload def __call__(self, inputs: tf.Tensor) -> tf.Tensor: ... @overload @@ -13,6 +14,9 @@ class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): @overload def __call__(self, inputs: tf.RaggedTensor) -> tf.RaggedTensor: ... # type: ignore def adapt( - self, data: tf.data.Dataset[TensorLike] | AnyArray, batch_size: int | None = None, steps: int | None = None + self, + data: tf.data.Dataset[TensorLike] | AnyArray | DataSequence, + batch_size: Integer | None = None, + steps: Float | None = None, ) -> None: ... - def compile(self, run_eagerly: bool | None = None, steps_per_execution: int | None = None) -> None: ... + def compile(self, run_eagerly: bool | None = None, steps_per_execution: Integer | None = None) -> None: ... From 31134a1b10d86b7cce9bd8f34713afdc925d3fd0 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 20:27:21 +0900 Subject: [PATCH 15/18] make IndexLookup private This class exists at runtime and is helpful for IntegerLookup/StringLookup, but it's undocumented internal class. --- .../tensorflow/keras/layers/preprocessing/__init__.pyi | 6 +++--- .../tensorflow/keras/layers/preprocessing/index_lookup.pyi | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi index 1896143f5a81..57214be16f7f 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi @@ -1,9 +1,9 @@ from typing import Literal from tensorflow._aliases import TensorCompatible -from tensorflow.keras.layers.preprocessing.index_lookup import IndexLookup +from tensorflow.keras.layers.preprocessing.index_lookup import _IndexLookup -class StringLookup(IndexLookup): +class StringLookup(_IndexLookup): def __init__( self, max_tokens: int | None = None, @@ -19,7 +19,7 @@ class StringLookup(IndexLookup): pad_to_max_tokens: bool = False, ) -> None: ... -class IntegerLookup(IndexLookup): +class IntegerLookup(_IndexLookup): def __init__( self, max_tokens: int | None = None, diff --git a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi index 3fd543a87e9a..8eaf0b627f34 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/index_lookup.pyi @@ -3,7 +3,7 @@ from _typeshed import Incomplete import tensorflow as tf from tensorflow.keras.layers.experimental.preprocessing import PreprocessingLayer -class IndexLookup(PreprocessingLayer): +class _IndexLookup(PreprocessingLayer): def compute_output_signature(self, input_spec: Incomplete) -> tf.TensorSpec: ... def get_vocabulary(self, include_special_tokens: bool = True) -> list[Incomplete]: ... def vocabulary_size(self) -> int: ... From aaa373868a1cc5235f380cbb8a680cadb8a82d55 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 20:38:52 +0900 Subject: [PATCH 16/18] silence/ignore mypy error. --- .../keras/layers/experimental/preprocessing.pyi | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi index 431222f18ecd..7fa074e39873 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi @@ -2,17 +2,21 @@ import abc from typing import overload import tensorflow as tf -from tensorflow._aliases import AnyArray, DataSequence, Float, Integer, TensorLike +from tensorflow._aliases import AnyArray, DataSequence, Float, Integer, TensorCompatible, TensorLike from tensorflow.keras.layers import Layer class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): is_adapted: bool + @overload # type: ignore + def __call__(self, inputs: tf.Tensor, *, training: bool = False, mask: TensorCompatible | None = None) -> tf.Tensor: ... @overload - def __call__(self, inputs: tf.Tensor) -> tf.Tensor: ... + def __call__( + self, inputs: tf.SparseTensor, *, training: bool = False, mask: TensorCompatible | None = None + ) -> tf.SparseTensor: ... @overload - def __call__(self, inputs: tf.SparseTensor) -> tf.SparseTensor: ... - @overload - def __call__(self, inputs: tf.RaggedTensor) -> tf.RaggedTensor: ... # type: ignore + def __call__( + self, inputs: tf.RaggedTensor, *, training: bool = False, mask: TensorCompatible | None = None + ) -> tf.RaggedTensor: ... def adapt( self, data: tf.data.Dataset[TensorLike] | AnyArray | DataSequence, From ac3b558e6b39db54bbf8e7ca733f956121e8fb7c Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 17 Feb 2024 21:07:25 +0900 Subject: [PATCH 17/18] fix: make PreprocessingLayer's is_adapted into a property. --- .../tensorflow/keras/layers/experimental/preprocessing.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi index 7fa074e39873..4d8cac5b812e 100644 --- a/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers/experimental/preprocessing.pyi @@ -6,7 +6,8 @@ from tensorflow._aliases import AnyArray, DataSequence, Float, Integer, TensorCo from tensorflow.keras.layers import Layer class PreprocessingLayer(Layer[TensorLike, TensorLike], metaclass=abc.ABCMeta): - is_adapted: bool + @property + def is_adapted(self) -> bool: ... @overload # type: ignore def __call__(self, inputs: tf.Tensor, *, training: bool = False, mask: TensorCompatible | None = None) -> tf.Tensor: ... @overload From 849279e13cf58ce29a59fa314ab44b012ba4fd45 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Wed, 13 Mar 2024 12:24:36 +0900 Subject: [PATCH 18/18] try to fix pytype issue Patch from this comment: https://github.com/python/typeshed/pull/11333#issuecomment-1993008007 --- stubs/tensorflow/tensorflow/__init__.pyi | 4 ---- stubs/tensorflow/tensorflow/_aliases.pyi | 4 +++- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 308690fe7556..a8436f1be77d 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -21,11 +21,7 @@ from tensorflow import ( from tensorflow._aliases import AnyArray, DTypeLike, ShapeLike, Slice, TensorCompatible from tensorflow.autodiff import GradientTape as GradientTape from tensorflow.core.protobuf import struct_pb2 - -# Explicit import of DType is covered by the wildcard, but -# is necessary to avoid a crash in pytype. from tensorflow.dtypes import * -from tensorflow.dtypes import DType as DType from tensorflow.experimental.dtensor import Layout from tensorflow.keras import losses as losses from tensorflow.linalg import eye as eye diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index 5566b80691d9..e26c80880309 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -5,9 +5,11 @@ from collections.abc import Iterable, Mapping, Sequence from typing import Any, Protocol, TypeVar from typing_extensions import TypeAlias +import numpy # pytype needs the unaliased import to resolve DTypeLike import numpy as np import numpy.typing as npt import tensorflow as tf +from tensorflow.dtypes import DType from tensorflow.keras.layers import InputSpec _T = TypeVar("_T") @@ -52,7 +54,7 @@ SparseTensorCompatible: TypeAlias = TensorCompatible | tf.SparseTensor TensorOrArray: TypeAlias = tf.Tensor | AnyArray ShapeLike: TypeAlias = tf.TensorShape | Iterable[ScalarTensorCompatible | None] | int | tf.Tensor -DTypeLike: TypeAlias = tf.DType | str | np.dtype[Any] | int +DTypeLike: TypeAlias = DType | str | numpy.dtype[Any] | int ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor] ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike]