adding alias types, wip

Author: hoel-bagard

stubs/tensorflow/tensorflow/__init__.pyi

@@ -19,7 +19,7 @@ from tensorflow import (
     keras as keras,
     math as math,
 )
-from tensorflow._aliases import ContainerGradients, ContainerTensors, ContainerTensorsLike, Gradients, TensorLike
+from tensorflow._aliases import _ContainerGradients, _ContainerTensors, _ContainerTensorsLike, _Gradients, _TensorLike
 from tensorflow.core.protobuf import struct_pb2
 
 # Explicit import of DType is covered by the wildcard, but
@@ -307,39 +307,39 @@ class GradientTape:
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
-        sources: TensorLike,
+        target: _ContainerTensors,
+        sources: _TensorLike,
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> Gradients: ...
+    ) -> _Gradients: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
+        target: _ContainerTensors,
         sources: Sequence[Tensor],
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> list[Gradients]: ...
+    ) -> list[_Gradients]: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
+        target: _ContainerTensors,
         sources: Mapping[str, Tensor],
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> dict[str, Gradients]: ...
+    ) -> dict[str, _Gradients]: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
-        sources: ContainerTensors,
+        target: _ContainerTensors,
+        sources: _ContainerTensors,
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> ContainerGradients: ...
+    ) -> _ContainerGradients: ...
     @contextmanager
     def stop_recording(self) -> Generator[None, None, None]: ...
     def reset(self) -> None: ...
-    def watch(self, tensor: ContainerTensorsLike) -> None: ...
+    def watch(self, tensor: _ContainerTensorsLike) -> None: ...
     def watched_variables(self) -> tuple[Variable, ...]: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 

stubs/tensorflow/tensorflow/_aliases.pyi

@@ -5,26 +5,53 @@
 from collections.abc import Mapping, Sequence
 from typing import Any, Protocol, TypeVar
 from typing_extensions import TypeAlias
+from typing import Any, Iterable, Mapping, Sequence, TypeVar
 
-import numpy
+import numpy as np
 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[_T1]] | Sequence[_ContainerGeneric[_T1]] | _T1
 
-TensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor | tf.SparseTensor
-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]
+_TensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor | tf.SparseTensor
+_SparseTensorLike = tf.Tensor | tf.SparseTensor
+_RaggedTensorLike = tf.Tensor | tf.RaggedTensor
+_RaggedTensorLikeT = TypeVar("_RaggedTensorLikeT", tf.Tensor, tf.RaggedTensor)
+_Gradients: TypeAlias = tf.Tensor | tf.IndexedSlices
 
 class _KerasSerializable1(Protocol):
     def get_config(self) -> dict[str, Any]: ...
 
 class _KerasSerializable2(Protocol):
     __name__: str
 
-KerasSerializable: TypeAlias = _KerasSerializable1 | _KerasSerializable2
+_KerasSerializable: TypeAlias = _KerasSerializable1 | _KerasSerializable2
+
+_FloatDataSequence = Sequence[float] | Sequence[_FloatDataSequence]
+_StrDataSequence = Sequence[str] | Sequence[_StrDataSequence]
+_ScalarTensorCompatible = tf.Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+
+_TensorCompatible = _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 = _TensorCompatible | tf.SparseTensor
+
+_ShapeLike = tf.TensorShape | Iterable[_ScalarTensorCompatible | None] | int | tf.Tensor
+_DTypeLike = tf.DType | str | np.dtype[Any] | int
+_GradientsT = tf.Tensor | tf.IndexedSlices
+
+_ContainerTensors = _ContainerGeneric[tf.Tensor]
+_ContainerTensorsLike = _ContainerGeneric[_TensorLike]
+_ContainerTensorCompatible = _ContainerGeneric[_TensorCompatible]
+_ContainerGradients = _ContainerGeneric[_GradientsT]
+_ContainerTensorShape = _ContainerGeneric[tf.TensorShape]
+_ContainerInputSpec = _ContainerGeneric[InputSpec]
+
+_AnyArray = np.ndarray[Any, Any]
+_FloatArray = np.ndarray[Any, np.dtype[np.float_ | np.float16 | np.float32 | np.float64]]
+_IntArray = np.ndarray[Any, np.dtype[np.int_ | np.uint8 | np.int32 | np.int64]]

stubs/tensorflow/tensorflow/data/__init__.pyi

@@ -7,7 +7,7 @@ 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._aliases import _ContainerGeneric
 from tensorflow.data import experimental as experimental
 from tensorflow.data.experimental import AUTOTUNE as AUTOTUNE
 from tensorflow.dtypes import DType
@@ -21,7 +21,7 @@ _T3 = TypeVar("_T3")
 class Iterator(_Iterator[_T1], Trackable, ABC):
     @property
     @abstractmethod
-    def element_spec(self) -> ContainerGeneric[TypeSpec[Any]]: ...
+    def element_spec(self) -> _ContainerGeneric[TypeSpec[Any]]: ...
     @abstractmethod
     def get_next(self) -> _T1: ...
     @abstractmethod
@@ -43,8 +43,8 @@ class Dataset(ABC, Generic[_T1]):
         element_length_func: Callable[[_T1], _ScalarTensorCompatible],
         bucket_boundaries: Sequence[int],
         bucket_batch_sizes: Sequence[int],
-        padded_shapes: ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
-        padding_values: ContainerGeneric[_ScalarTensorCompatible] | None = None,
+        padded_shapes: _ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
+        padding_values: _ContainerGeneric[_ScalarTensorCompatible] | None = None,
         pad_to_bucket_boundary: bool = False,
         no_padding: bool = False,
         drop_remainder: bool = False,
@@ -63,18 +63,18 @@ class Dataset(ABC, Generic[_T1]):
     ) -> Dataset[tf.Tensor]: ...
     @property
     @abstractmethod
-    def element_spec(self) -> ContainerGeneric[TypeSpec[Any]]: ...
+    def element_spec(self) -> _ContainerGeneric[TypeSpec[Any]]: ...
     def enumerate(self, start: _ScalarTensorCompatible = 0, name: str | None = None) -> Dataset[tuple[int, _T1]]: ...
     def filter(self, predicate: Callable[[_T1], bool | tf.Tensor], name: str | None = None) -> Dataset[_T1]: ...
     def flat_map(self, map_func: Callable[[_T1], Dataset[_T2]], name: str | None = None) -> Dataset[_T2]: ...
     # PEP 646 can be used here for a more precise type when better supported.
     @staticmethod
     def from_generator(
         generator: Callable[..., _T2],
-        output_types: ContainerGeneric[DType] | None = None,
-        output_shapes: ContainerGeneric[tf.TensorShape | Sequence[int | None]] | None = None,
+        output_types: _ContainerGeneric[DType] | None = None,
+        output_shapes: _ContainerGeneric[tf.TensorShape | Sequence[int | None]] | None = None,
         args: tuple[object, ...] | None = None,
-        output_signature: ContainerGeneric[TypeSpec[Any]] | None = None,
+        output_signature: _ContainerGeneric[TypeSpec[Any]] | None = None,
         name: str | None = None,
     ) -> Dataset[_T2]: ...
     @staticmethod
@@ -111,7 +111,7 @@ class Dataset(ABC, Generic[_T1]):
     @staticmethod
     def load(
         path: str,
-        element_spec: ContainerGeneric[tf.TypeSpec[Any]] | None = None,
+        element_spec: _ContainerGeneric[tf.TypeSpec[Any]] | None = None,
         compression: _CompressionTypes = None,
         reader_func: Callable[[Dataset[Dataset[Any]]], Dataset[Any]] | None = None,
     ) -> Dataset[Any]: ...
@@ -127,8 +127,8 @@ class Dataset(ABC, Generic[_T1]):
     def padded_batch(
         self,
         batch_size: _ScalarTensorCompatible,
-        padded_shapes: ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
-        padding_values: ContainerGeneric[_ScalarTensorCompatible] | None = None,
+        padded_shapes: _ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
+        padding_values: _ContainerGeneric[_ScalarTensorCompatible] | None = None,
         drop_remainder: bool = False,
         name: str | None = None,
     ) -> Dataset[_T1]: ...

stubs/tensorflow/tensorflow/io/__init__.pyi

@@ -5,7 +5,7 @@ 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 _TensorLike
 from tensorflow.io import gfile as gfile
 
 _FeatureSpecs: TypeAlias = Mapping[str, FixedLenFeature | FixedLenSequenceFeature | VarLenFeature | RaggedFeature | SparseFeature]
@@ -102,5 +102,5 @@ class RaggedFeature(NamedTuple):
 
 def parse_example(
     serialized: _TensorCompatible, features: _FeatureSpecs, example_names: Iterable[str] | None = None, name: str | None = None
-) -> dict[str, TensorLike]: ...
+) -> dict[str, _TensorLike]: ...
 def __getattr__(name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/keras/__init__.pyi

@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from typing import Callable
 
 from tensorflow.keras import (
     activations as activations,
@@ -12,5 +13,20 @@ from tensorflow.keras import (
     regularizers as regularizers,
 )
 from tensorflow.keras.models import Model as Model
+import tensorflow as tf
+from tensorflow._aliases import _TensorCompatible
 
 def __getattr__(name: str) -> Incomplete: ...
+
+_Loss = (
+   str
+    | tf.keras.losses.Loss
+    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
+)
+
+_Metric = (
+    str
+    | tf.keras.metrics.Metric
+    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
+    | None
+)

stubs/tensorflow/tensorflow/keras/layers.pyi

@@ -5,7 +5,7 @@ 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._aliases import _AnyArray
 from tensorflow.keras.activations import _Activation
 from tensorflow.keras.constraints import Constraint
 from tensorflow.keras.initializers import _Initializer
@@ -90,8 +90,8 @@ class Layer(tf.Module, Generic[_InputT, _OutputT]):
     def non_trainable_weights(self) -> list[Variable]: ...
     @property
     def losses(self) -> list[Tensor]: ...
-    def get_weights(self) -> list[AnyArray]: ...
-    def set_weights(self, weights: Sequence[AnyArray]) -> None: ...
+    def get_weights(self) -> list[_AnyArray]: ...
+    def set_weights(self, weights: Sequence[_AnyArray]) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> Self: ...

stubs/tensorflow/tensorflow/keras/losses.pyi

@@ -5,7 +5,7 @@ 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._aliases import _KerasSerializable
 from tensorflow.keras.metrics import (
     binary_crossentropy as binary_crossentropy,
     categorical_crossentropy as categorical_crossentropy,
@@ -136,7 +136,7 @@ def log_cosh(y_true: _TensorCompatible, y_pred: _TensorCompatible) -> Tensor: ..
 def deserialize(
     name: str | dict[str, Any], custom_objects: dict[str, Any] | None = None, use_legacy_format: bool = False
 ) -> Loss: ...
-def serialize(loss: KerasSerializable, use_legacy_format: bool = False) -> dict[str, Any]: ...
+def serialize(loss: _KerasSerializable, use_legacy_format: bool = False) -> dict[str, Any]: ...
 
 _FuncT = TypeVar("_FuncT", bound=Callable[..., Any])
 

stubs/tensorflow/tensorflow/keras/models.pyi

@@ -10,9 +10,18 @@ import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
-from tensorflow.keras.optimizers.legacy import Optimizer
+from tensorflow.keras import _Loss, _Metric
+from tensorflow._aliases import _ContainerGeneric
+
+_BothOptimizer = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
+    _train_counter: tf.Variable
+    _test_counter: tf.Variable
+    optimizer: _BothOptimizer | None
+    loss: tf.keras.losses.Loss | dict[str, tf.keras.losses.Loss]
+    stop_training: bool
+
     def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def __setattr__(self, name: str, value: Any) -> None: ...
@@ -21,13 +30,14 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    # Ideally loss/metrics/output would share the same structure but higher kinded types are not supported.
     def compile(
         self,
-        optimizer: Optimizer | str = "rmsprop",
-        loss: tf.keras.losses.Loss | str | None = None,
-        metrics: list[tf.keras.metrics.Metric | str] | None = None,
-        loss_weights: list[float] | dict[str, float] | None = None,
-        weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
+        optimizer: _BothOptimizer | str = "rmsprop",
+        loss: _ContainerGeneric[_Loss] | None = None,
+        metrics: _ContainerGeneric[_Metric] | None = None,
+        loss_weights: _ContainerGeneric[float] | None = None,
+        weighted_metrics: _ContainerGeneric[_Metric] | None = None,
         run_eagerly: bool | None = None,
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,

stubs/tensorflow/tensorflow/keras/optimizers/legacy/__init__.pyi

@@ -5,17 +5,17 @@ from typing import Any
 from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
-from tensorflow._aliases import Gradients
+from tensorflow._aliases import _Gradients
 from tensorflow.keras.optimizers import schedules as schedules
 from tensorflow.python.trackable.base import Trackable
 
 _Initializer: TypeAlias = str | Callable[[], tf.Tensor] | dict[str, Any]
 _Shape: TypeAlias = tf.TensorShape | Iterable[int | None]
 _Dtype: TypeAlias = tf.DType | str | None
 _LearningRate: TypeAlias = float | tf.Tensor | schedules.LearningRateSchedule | Callable[[], float | tf.Tensor]
-_GradientAggregator: TypeAlias = Callable[[list[tuple[Gradients, tf.Variable]]], list[tuple[Gradients, tf.Variable]]] | None
+_GradientAggregator: TypeAlias = Callable[[list[tuple[_Gradients, tf.Variable]]], list[tuple[_Gradients, tf.Variable]]] | None
 _GradientTransformer: TypeAlias = (
-    Iterable[Callable[[list[tuple[Gradients, tf.Variable]]], list[tuple[Gradients, tf.Variable]]]] | None
+    Iterable[Callable[[list[tuple[_Gradients, tf.Variable]]], list[tuple[_Gradients, tf.Variable]]]] | None
 )
 
 # kwargs here and in other optimizers can be given better type after Unpack[TypedDict], PEP 692, is supported.
@@ -53,7 +53,7 @@ class Optimizer(Trackable):
     ) -> tf.Variable: ...
     def apply_gradients(
         self,
-        grads_and_vars: Iterable[tuple[Gradients, tf.Variable]],
+        grads_and_vars: Iterable[tuple[_Gradients, tf.Variable]],
         name: str | None = None,
         experimental_aggregate_gradients: bool = True,
     ) -> tf.Operation | None: ...
@@ -64,7 +64,7 @@ class Optimizer(Trackable):
     def get_config(self) -> dict[str, Any]: ...
     def get_slot(self, var: tf.Variable, slot_name: str) -> tf.Variable: ...
     def get_slot_names(self) -> list[str]: ...
-    def get_gradients(self, loss: tf.Tensor, params: list[tf.Variable]) -> list[Gradients]: ...
+    def get_gradients(self, loss: tf.Tensor, params: list[tf.Variable]) -> list[_Gradients]: ...
     def minimize(
         self,
         loss: tf.Tensor | Callable[[], tf.Tensor],