fix imports.

Author: hoel-bagard

stubs/tensorflow/tensorflow/data/__init__.pyi

@@ -6,8 +6,8 @@ from typing_extensions import Self
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import TypeSpec, _ScalarTensorCompatible, _TensorCompatible
-from tensorflow._aliases import _ContainerGeneric
+from tensorflow import TypeSpec
+from tensorflow._aliases import _ContainerGeneric, _ScalarTensorCompatible, _TensorCompatible
 from tensorflow.data import experimental as experimental
 from tensorflow.data.experimental import AUTOTUNE as AUTOTUNE
 from tensorflow.dtypes import DType

stubs/tensorflow/tensorflow/data/experimental.pyi

@@ -2,7 +2,8 @@ from _typeshed import Incomplete
 from collections.abc import Callable, Sequence
 from typing import Final, TypeVar
 
-from tensorflow import Tensor, _TensorCompatible
+from tensorflow import Tensor
+from tensorflow._aliases import _TensorCompatible
 from tensorflow.data import Dataset
 
 AUTOTUNE: Final = -1

stubs/tensorflow/tensorflow/dtypes.pyi

@@ -4,7 +4,7 @@ from builtins import bool as _bool
 from typing import Any
 
 import numpy as np
-from tensorflow import _DTypeLike
+from tensorflow._aliases import _DTypeLike
 
 class _DTypeMeta(ABCMeta): ...
 

stubs/tensorflow/tensorflow/feature_column/__init__.pyi

@@ -1,7 +1,7 @@
 from collections.abc import Callable, Iterable, Sequence
 
 import tensorflow as tf
-from tensorflow import _ShapeLike
+from tensorflow._aliases import _ShapeLike
 from tensorflow.python.feature_column import feature_column_v2 as fc, sequence_feature_column as seq_fc
 
 def numeric_column(

stubs/tensorflow/tensorflow/io/__init__.pyi

@@ -4,8 +4,7 @@ from types import TracebackType
 from typing import Literal, NamedTuple
 from typing_extensions import Self, TypeAlias
 
-from tensorflow import _DTypeLike, _ShapeLike, _TensorCompatible
-from tensorflow._aliases import _TensorLike
+from tensorflow._aliases import _DTypeLike, _ShapeLike, _TensorCompatible, _TensorLike
 from tensorflow.io import gfile as gfile
 
 _FeatureSpecs: TypeAlias = Mapping[str, FixedLenFeature | FixedLenSequenceFeature | VarLenFeature | RaggedFeature | SparseFeature]

stubs/tensorflow/tensorflow/keras/initializers.pyi

@@ -3,7 +3,8 @@ from collections.abc import Callable
 from typing import Any, overload
 from typing_extensions import Self, TypeAlias
 
-from tensorflow import Tensor, _DTypeLike, _ShapeLike, _TensorCompatible
+from tensorflow import Tensor
+from tensorflow._aliases import _DTypeLike, _ShapeLike, _TensorCompatible
 
 class Initializer:
     def __call__(self, shape: _ShapeLike, dtype: _DTypeLike | None = None) -> Tensor: ...

stubs/tensorflow/tensorflow/keras/layers.pyi

@@ -4,8 +4,8 @@ from typing import Any, Generic, TypeVar, overload
 from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
-from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible
-from tensorflow._aliases import _AnyArray
+from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization
+from tensorflow._aliases import _AnyArray, _DTypeLike, _TensorCompatible
 from tensorflow.keras.activations import _Activation
 from tensorflow.keras.constraints import Constraint
 from tensorflow.keras.initializers import _Initializer
@@ -23,7 +23,7 @@ class InputSpec:
     axes: dict[int, int | None] | None
     def __init__(
         self,
-        dtype: tf._DTypeLike | None = None,
+        dtype: _DTypeLike | None = None,
         shape: Iterable[int | None] | None = None,
         ndim: int | None = None,
         max_ndim: int | None = None,
@@ -49,7 +49,7 @@ class Layer(tf.Module, Generic[_InputT, _OutputT]):
     @trainable.setter
     def trainable(self, value: bool) -> None: ...
     def __init__(
-        self, trainable: bool = True, name: str | None = None, dtype: tf._DTypeLike | None = None, dynamic: bool = False
+        self, trainable: bool = True, name: str | None = None, dtype: _DTypeLike | None = None, dynamic: bool = False
     ) -> None: ...
 
     # *args/**kwargs are allowed, but have obscure footguns and tensorflow documentation discourages their usage.
@@ -69,7 +69,7 @@ class Layer(tf.Module, Generic[_InputT, _OutputT]):
         self,
         name: str | None = None,
         shape: Iterable[int | None] | None = None,
-        dtype: tf._DTypeLike | None = None,
+        dtype: _DTypeLike | None = None,
         initializer: _Initializer | None = None,
         regularizer: _Regularizer = None,
         trainable: bool | None = None,
@@ -106,7 +106,7 @@ class Layer(tf.Module, Generic[_InputT, _OutputT]):
 # all layer constructors.
 
 # TODO: Replace last Any after adding tf.keras.mixed_precision.Policy.
-_LayerDtype: TypeAlias = tf._DTypeLike | dict[str, Any] | Any
+_LayerDtype: TypeAlias = _DTypeLike | dict[str, Any] | Any
 
 _Constraint: TypeAlias = str | dict[str, Any] | Constraint | None
 

stubs/tensorflow/tensorflow/keras/losses.pyi

@@ -4,8 +4,8 @@ from collections.abc import Callable
 from typing import Any, Final, Literal, TypeVar, overload
 from typing_extensions import Self, TypeAlias, TypeGuard
 
-from tensorflow import Tensor, _TensorCompatible
-from tensorflow._aliases import _KerasSerializable
+from tensorflow import Tensor
+from tensorflow._aliases import _KerasSerializable, _TensorCompatible
 from tensorflow.keras.metrics import (
     binary_crossentropy as binary_crossentropy,
     categorical_crossentropy as categorical_crossentropy,

stubs/tensorflow/tensorflow/keras/metrics.pyi

@@ -1,4 +1,5 @@
-from tensorflow import Tensor, _TensorCompatible
+from tensorflow import Tensor
+from tensorflow._aliases import _TensorCompatible
 
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1

stubs/tensorflow/tensorflow/math.pyi

@@ -3,7 +3,8 @@ from collections.abc import Iterable
 from typing import TypeVar, overload
 from typing_extensions import TypeAlias
 
-from tensorflow import IndexedSlices, RaggedTensor, Tensor, _DTypeLike, _ShapeLike, _TensorCompatible
+from tensorflow import IndexedSlices, RaggedTensor, Tensor
+from tensorflow._aliases import _DTypeLike, _ShapeLike, _TensorCompatible
 from tensorflow.sparse import SparseTensor
 
 _TensorCompatibleT = TypeVar("_TensorCompatibleT", bound=_TensorCompatible)