fix: fix imports

python · Jan 31, 2024 · 5ee5f57 · 5ee5f57
1 parent 518b69e
commit 5ee5f57
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi b/stubs/tensorflow/tensorflow/keras/layers/preprocessing/__init__.pyi
@@ -1,6 +1,6 @@
 from typing import Literal
 
-from tensorflow import _TensorCompatible
+from tensorflow._aliases import TensorCompatible
 from tensorflow.keras.layers.preprocessing.index_lookup import IndexLookup
 
 class StringLookup(IndexLookup):
@@ -10,8 +10,8 @@ class StringLookup(IndexLookup):
         num_oov_indices: int = 1,
         mask_token: str | None = None,
         oov_token: str = "[UNK]",
-        vocabulary: str | None | _TensorCompatible = None,
-        idf_weights: _TensorCompatible | None = None,
+        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",
@@ -26,9 +26,9 @@ class IntegerLookup(IndexLookup):
         num_oov_indices: int = 1,
         mask_token: int | None = None,
         oov_token: int = -1,
-        vocabulary: str | None | _TensorCompatible = None,
+        vocabulary: str | None | TensorCompatible = None,
         vocabulary_dtype: Literal["int64", "int32"] = "int64",
-        idf_weights: _TensorCompatible | None = None,
+        idf_weights: TensorCompatible | None = None,
         invert: bool = False,
         output_mode: Literal["int", "count", "multi_hot", "one_hot", "tf_idf"] = "int",
         sparse: bool = False,