From 22c88bc190694a08f5a8dc2a6c7ba1b445248068 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 28 Jan 2024 18:34:30 +0900
Subject: [PATCH] Move some layers to keras.layers.preprocessing

[pre-commit.ci] auto fixes from pre-commit.com hooks
---
 .../keras/{layers.pyi => layers/__init__.pyi} | 42 +------------------
 .../layers/experimental/preprocessing.pyi     |  6 +++
 .../keras/layers/preprocessing/__init__.pyi   | 36 ++++++++++++++++
 .../layers/preprocessing/index_lookup.pyi     |  9 ++++
 4 files changed, 52 insertions(+), 41 deletions(-)
 rename stubs/tensorflow/tensorflow/keras/{layers.pyi => layers/__init__.pyi} (88%)
 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 88%
rename from stubs/tensorflow/tensorflow/keras/layers.pyi
rename to stubs/tensorflow/tensorflow/keras/layers/__init__.pyi
index 8da6daefb67f..6b3f192c1e2e 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
 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
 
@@ -244,47 +245,6 @@ class LayerNormalization(Layer[tf.Tensor, tf.Tensor]):
         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,
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..9dae75b8397c
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi
@@ -0,0 +1,36 @@
+from typing import Literal
+
+from tensorflow 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: ...