tensorflow: Add (and rename) aliases

stubs/tensorflow/tensorflow/__init__.pyi

@@ -6,9 +6,8 @@ from contextlib import contextmanager
 from enum import Enum
 from types import TracebackType
 from typing import Any, Generic, NoReturn, TypeVar, overload
-from typing_extensions import ParamSpec, Self, TypeAlias
+from typing_extensions import ParamSpec, Self
 
-import numpy
 from google.protobuf.message import Message
 from tensorflow import (
     data as data,
@@ -19,7 +18,18 @@ from tensorflow import (
     keras as keras,
     math as math,
 )
-from tensorflow._aliases import ContainerGradients, ContainerTensors, ContainerTensorsLike, Gradients, TensorLike
+from tensorflow._aliases import (
+    AnyArray,
+    ContainerGradients,
+    ContainerTensors,
+    ContainerTensorsLike,
+    DTypeLike,
+    Gradients,
+    ShapeLike,
+    Slice,
+    TensorCompatible,
+    TensorLike,
+)
 from tensorflow.core.protobuf import struct_pb2
 
 # Explicit import of DType is covered by the wildcard, but
@@ -73,15 +83,6 @@ from tensorflow.sparse import SparseTensor as SparseTensor
 # we will skip making Tensor generic. Also good type hints for shapes will
 # run quickly into many places where type system is not strong enough today.
 # So shape typing is probably not worth doing anytime soon.
-_Slice: TypeAlias = int | slice | None
-
-_FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
-_StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
-_ScalarTensorCompatible: TypeAlias = Tensor | str | float | numpy.ndarray[Any, Any] | numpy.number[Any]
-_TensorCompatible: TypeAlias = _ScalarTensorCompatible | Sequence[_TensorCompatible]
-_ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
-_DTypeLike: TypeAlias = DType | str | numpy.dtype[Any]
-
 class Tensor:
     def __init__(self, op: Operation, value_index: int, dtype: DType) -> None: ...
     def consumers(self) -> list[Incomplete]: ...
@@ -96,35 +97,35 @@ class Tensor:
     def name(self) -> str: ...
     @property
     def op(self) -> Operation: ...
-    def numpy(self) -> numpy.ndarray[Any, Any]: ...
+    def numpy(self) -> AnyArray: ...
     def __int__(self) -> int: ...
     def __abs__(self, name: str | None = None) -> Tensor: ...
-    def __add__(self, other: _TensorCompatible) -> Tensor: ...
-    def __radd__(self, other: _TensorCompatible) -> Tensor: ...
-    def __sub__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rsub__(self, other: _TensorCompatible) -> Tensor: ...
-    def __mul__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rmul__(self, other: _TensorCompatible) -> Tensor: ...
-    def __pow__(self, other: _TensorCompatible) -> Tensor: ...
-    def __matmul__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rmatmul__(self, other: _TensorCompatible) -> Tensor: ...
-    def __floordiv__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rfloordiv__(self, other: _TensorCompatible) -> Tensor: ...
-    def __truediv__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rtruediv__(self, other: _TensorCompatible) -> Tensor: ...
+    def __add__(self, other: TensorCompatible) -> Tensor: ...
+    def __radd__(self, other: TensorCompatible) -> Tensor: ...
+    def __sub__(self, other: TensorCompatible) -> Tensor: ...
+    def __rsub__(self, other: TensorCompatible) -> Tensor: ...
+    def __mul__(self, other: TensorCompatible) -> Tensor: ...
+    def __rmul__(self, other: TensorCompatible) -> Tensor: ...
+    def __pow__(self, other: TensorCompatible) -> Tensor: ...
+    def __matmul__(self, other: TensorCompatible) -> Tensor: ...
+    def __rmatmul__(self, other: TensorCompatible) -> Tensor: ...
+    def __floordiv__(self, other: TensorCompatible) -> Tensor: ...
+    def __rfloordiv__(self, other: TensorCompatible) -> Tensor: ...
+    def __truediv__(self, other: TensorCompatible) -> Tensor: ...
+    def __rtruediv__(self, other: TensorCompatible) -> Tensor: ...
     def __neg__(self, name: str | None = None) -> Tensor: ...
-    def __and__(self, other: _TensorCompatible) -> Tensor: ...
-    def __rand__(self, other: _TensorCompatible) -> Tensor: ...
-    def __or__(self, other: _TensorCompatible) -> Tensor: ...
-    def __ror__(self, other: _TensorCompatible) -> Tensor: ...
-    def __eq__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
-    def __ne__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
-    def __ge__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
-    def __gt__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
-    def __le__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
-    def __lt__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __and__(self, other: TensorCompatible) -> Tensor: ...
+    def __rand__(self, other: TensorCompatible) -> Tensor: ...
+    def __or__(self, other: TensorCompatible) -> Tensor: ...
+    def __ror__(self, other: TensorCompatible) -> Tensor: ...
+    def __eq__(self, other: TensorCompatible) -> Tensor: ...  # type: ignore[override]
+    def __ne__(self, other: TensorCompatible) -> Tensor: ...  # type: ignore[override]
+    def __ge__(self, other: TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __gt__(self, other: TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __le__(self, other: TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __lt__(self, other: TensorCompatible, name: str | None = None) -> Tensor: ...
     def __bool__(self) -> NoReturn: ...
-    def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> Tensor: ...
+    def __getitem__(self, slice_spec: Slice | tuple[Slice, ...]) -> Tensor: ...
     def __len__(self) -> int: ...
     # This only works for rank 0 tensors.
     def __index__(self) -> int: ...
@@ -160,25 +161,25 @@ class Variable(Tensor, metaclass=_VariableMetaclass):
         # Real type is VariableDef protobuf type. Can be added after adding script
         # to generate tensorflow protobuf stubs with mypy-protobuf.
         variable_def: Incomplete | None = None,
-        dtype: _DTypeLike | None = None,
+        dtype: DTypeLike | None = None,
         import_scope: str | None = None,
         constraint: Callable[[Tensor], Tensor] | None = None,
         synchronization: VariableSynchronization = ...,
         aggregation: VariableAggregation = ...,
-        shape: _ShapeLike | None = None,
+        shape: ShapeLike | None = None,
         experimental_enable_variable_lifting: _bool = True,
     ) -> None: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class RaggedTensor(metaclass=ABCMeta):
     def bounding_shape(
-        self, axis: _TensorCompatible | None = None, name: str | None = None, out_type: _DTypeLike | None = None
+        self, axis: TensorCompatible | None = None, name: str | None = None, out_type: DTypeLike | None = None
     ) -> Tensor: ...
     @classmethod
-    def from_sparse(cls, st_input: SparseTensor, name: str | None = None, row_splits_dtype: _DTypeLike = ...) -> RaggedTensor: ...
+    def from_sparse(cls, st_input: SparseTensor, name: str | None = None, row_splits_dtype: DTypeLike = ...) -> RaggedTensor: ...
     def to_sparse(self, name: str | None = None) -> SparseTensor: ...
     def to_tensor(
-        self, default_value: float | str | None = None, name: str | None = None, shape: _ShapeLike | None = None
+        self, default_value: float | str | None = None, name: str | None = None, shape: ShapeLike | None = None
     ) -> Tensor: ...
     def __add__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
     def __radd__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
@@ -187,7 +188,7 @@ class RaggedTensor(metaclass=ABCMeta):
     def __rmul__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
     def __floordiv__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
     def __truediv__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
-    def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> RaggedTensor: ...
+    def __getitem__(self, slice_spec: Slice | tuple[Slice, ...]) -> RaggedTensor: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class Operation:
@@ -216,7 +217,7 @@ class Operation:
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class TensorShape(metaclass=ABCMeta):
-    def __init__(self, dims: _ShapeLike) -> None: ...
+    def __init__(self, dims: ShapeLike) -> None: ...
     @property
     def rank(self) -> int: ...
     def as_list(self) -> list[int | None]: ...
@@ -354,14 +355,14 @@ class TypeSpec(ABC, Generic[_SpecProto]):
     def experimental_from_proto(cls, proto: _SpecProto) -> Self: ...
     @classmethod
     def experimental_type_proto(cls) -> type[_SpecProto]: ...
-    def is_compatible_with(self, spec_or_value: Self | _TensorCompatible | SparseTensor | RaggedTensor) -> _bool: ...
+    def is_compatible_with(self, spec_or_value: Self | TensorCompatible | SparseTensor | RaggedTensor) -> _bool: ...
     # Incomplete as tf.types is not yet covered.
     def is_subtype_of(self, other: Incomplete) -> _bool: ...
     def most_specific_common_supertype(self, others: Sequence[Incomplete]) -> Self | None: ...
     def most_specific_compatible_type(self, other: Self) -> Self: ...
 
 class TensorSpec(TypeSpec[struct_pb2.TensorSpecProto]):
-    def __init__(self, shape: _ShapeLike, dtype: _DTypeLike = ..., name: str | None = None) -> None: ...
+    def __init__(self, shape: ShapeLike, dtype: DTypeLike = ..., name: str | None = None) -> None: ...
     @property
     def value_type(self) -> Tensor: ...
     @property
@@ -374,10 +375,10 @@ class TensorSpec(TypeSpec[struct_pb2.TensorSpecProto]):
     def from_spec(cls, spec: TypeSpec[Any], name: str | None = None) -> Self: ...
     @classmethod
     def from_tensor(cls, tensor: Tensor, name: str | None = None) -> Self: ...
-    def is_compatible_with(self, spec_or_tensor: Self | _TensorCompatible) -> _bool: ...  # type: ignore[override]
+    def is_compatible_with(self, spec_or_tensor: Self | TensorCompatible) -> _bool: ...  # type: ignore[override]
 
 class SparseTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
-    def __init__(self, shape: _ShapeLike | None = None, dtype: _DTypeLike = ...) -> None: ...
+    def __init__(self, shape: ShapeLike | None = None, dtype: DTypeLike = ...) -> None: ...
     @property
     def value_type(self) -> SparseTensor: ...
     @property
@@ -390,10 +391,10 @@ class SparseTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
 class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
     def __init__(
         self,
-        shape: _ShapeLike | None = None,
-        dtype: _DTypeLike = ...,
+        shape: ShapeLike | None = None,
+        dtype: DTypeLike = ...,
         ragged_rank: int | None = None,
-        row_splits_dtype: _DTypeLike = ...,
+        row_splits_dtype: DTypeLike = ...,
         flat_values_spec: TypeSpec[Any] | None = None,
     ) -> None: ...
     @property
@@ -407,35 +408,35 @@ class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
 
 def __getattr__(name: str) -> Incomplete: ...
 def convert_to_tensor(
-    value: _TensorCompatible | IndexedSlices,
-    dtype: _DTypeLike | None = None,
-    dtype_hint: _DTypeLike | None = None,
+    value: TensorCompatible | IndexedSlices,
+    dtype: DTypeLike | None = None,
+    dtype_hint: DTypeLike | None = None,
     name: str | None = None,
 ) -> Tensor: ...
 @overload
-def expand_dims(input: _TensorCompatible, axis: int, name: str | None = None) -> Tensor: ...
+def expand_dims(input: TensorCompatible, axis: int, name: str | None = None) -> Tensor: ...
 @overload
 def expand_dims(input: RaggedTensor, axis: int, name: str | None = None) -> RaggedTensor: ...
 @overload
-def concat(values: _TensorCompatible, axis: int, name: str | None = "concat") -> Tensor: ...
+def concat(values: TensorCompatible, axis: int, name: str | None = "concat") -> Tensor: ...
 @overload
 def concat(values: Sequence[RaggedTensor], axis: int, name: str | None = "concat") -> RaggedTensor: ...
 @overload
 def squeeze(
-    input: _TensorCompatible, axis: int | tuple[int, ...] | list[int] | None = None, name: str | None = None
+    input: TensorCompatible, axis: int | tuple[int, ...] | list[int] | None = None, name: str | None = None
 ) -> Tensor: ...
 @overload
 def squeeze(input: RaggedTensor, axis: int | tuple[int, ...] | list[int], name: str | None = None) -> RaggedTensor: ...
 def tensor_scatter_nd_update(
-    tensor: _TensorCompatible, indices: _TensorCompatible, updates: _TensorCompatible, name: str | None = None
+    tensor: TensorCompatible, indices: TensorCompatible, updates: TensorCompatible, name: str | None = None
 ) -> Tensor: ...
 def constant(
-    value: _TensorCompatible, dtype: _DTypeLike | None = None, shape: _ShapeLike | None = None, name: str | None = "Const"
+    value: TensorCompatible, dtype: DTypeLike | None = None, shape: ShapeLike | None = None, name: str | None = "Const"
 ) -> Tensor: ...
 @overload
-def cast(x: _TensorCompatible, dtype: _DTypeLike, name: str | None = None) -> Tensor: ...
+def cast(x: TensorCompatible, dtype: DTypeLike, name: str | None = None) -> Tensor: ...
 @overload
-def cast(x: SparseTensor, dtype: _DTypeLike, name: str | None = None) -> SparseTensor: ...
+def cast(x: SparseTensor, dtype: DTypeLike, name: str | None = None) -> SparseTensor: ...
 @overload
-def cast(x: RaggedTensor, dtype: _DTypeLike, name: str | None = None) -> RaggedTensor: ...
-def reshape(tensor: _TensorCompatible, shape: _ShapeLike | Tensor, name: str | None = None) -> Tensor: ...
+def cast(x: RaggedTensor, dtype: DTypeLike, name: str | None = None) -> RaggedTensor: ...
+def reshape(tensor: TensorCompatible, shape: ShapeLike | Tensor, name: str | None = None) -> Tensor: ...

stubs/tensorflow/tensorflow/_aliases.pyi

@@ -1,30 +1,55 @@
 # Commonly used type aliases.
-# Everything in this module is private for stubs. There is no runtime
-# equivalent.
+# Everything in this module is private for stubs. There is no runtime equivalent.
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Protocol, TypeVar
 from typing_extensions import TypeAlias
 
-import numpy
+import numpy as np
+import numpy.typing as npt
 import tensorflow as tf
+from tensorflow.keras.layers import InputSpec
 
-_T1 = TypeVar("_T1")
-ContainerGeneric: TypeAlias = Mapping[str, ContainerGeneric[_T1]] | Sequence[ContainerGeneric[_T1]] | _T1
+_T = TypeVar("_T")
+ContainerGeneric: TypeAlias = Mapping[str, ContainerGeneric[_T]] | Sequence[ContainerGeneric[_T]] | _T
 
 TensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor | tf.SparseTensor
+SparseTensorLike: TypeAlias = tf.Tensor | tf.SparseTensor
+RaggedTensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor
+# _RaggedTensorLikeT = TypeVar("_RaggedTensorLikeT", tf.Tensor, tf.RaggedTensor)
 Gradients: TypeAlias = tf.Tensor | tf.IndexedSlices
 
-ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike]
-ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor]
-ContainerGradients: TypeAlias = ContainerGeneric[Gradients]
-
-AnyArray: TypeAlias = numpy.ndarray[Any, Any]
-
-class _KerasSerializable1(Protocol):
+class KerasSerializable1(Protocol):
     def get_config(self) -> dict[str, Any]: ...
 
-class _KerasSerializable2(Protocol):
+class KerasSerializable2(Protocol):
     __name__: str
 
-KerasSerializable: TypeAlias = _KerasSerializable1 | _KerasSerializable2
+KerasSerializable: TypeAlias = KerasSerializable1 | KerasSerializable2
+
+Slice: TypeAlias = int | slice | None
+FloatDataSequence: TypeAlias = Sequence[float] | Sequence[FloatDataSequence]
+StrDataSequence: TypeAlias = Sequence[str] | Sequence[StrDataSequence]
+ScalarTensorCompatible: TypeAlias = tf.Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+
+TensorCompatible: TypeAlias = ScalarTensorCompatible | Sequence[TensorCompatible]
+# _TensorCompatibleT = TypeVar("_TensorCompatibleT", bound=TensorCompatible)
+# Sparse tensors are very annoying. Some operations work on them, but many do not.
+# You will need to manually verify if an operation supports them. SparseTensorCompatible is intended to be a
+# broader type than TensorCompatible and not all operations will support broader version. If unsure,
+# use TensorCompatible instead.
+SparseTensorCompatible: TypeAlias = TensorCompatible | tf.SparseTensor
+
+ShapeLike: TypeAlias = tf.TensorShape | Iterable[ScalarTensorCompatible | None] | int | tf.Tensor
+DTypeLike: TypeAlias = tf.DType | str | np.dtype[Any] | int
+
+ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor]
+ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike]
+ContainerTensorCompatible: TypeAlias = ContainerGeneric[TensorCompatible]
+ContainerGradients: TypeAlias = ContainerGeneric[Gradients]
+ContainerTensorShape: TypeAlias = ContainerGeneric[tf.TensorShape]
+ContainerInputSpec: TypeAlias = ContainerGeneric[InputSpec]
+
+AnyArray: TypeAlias = npt.NDArray[Any]
+FloatArray: TypeAlias = npt.NDArray[np.float_ | np.float16 | np.float32 | np.float64]
+IntArray: TypeAlias = npt.NDArray[np.int_ | np.uint8 | np.int32 | np.int64]