Add TensorFlow stubs

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
@@ -73,7 +73,7 @@ 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
+_Slice: TypeAlias = int | slice | Tensor | None
 
 _FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
 _StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
@@ -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: ...
 
@@ -407,8 +407,14 @@ 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,
-    name: str | None = None,
+    value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
+def expand_dims(input: Tensor | tuple[int, ...], axis: int, name: None | str = None) -> Tensor: ...
+def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
+def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
+def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
+def constant(
+    value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = "Const"
+) -> Tensor: ...
+def cast(x: Tensor | _TensorCompatible, dtype: DType, name: None | str = None) -> Tensor: ...
+def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...

stubs/tensorflow/tensorflow/_aliases.pyi

@@ -3,28 +3,54 @@
 # equivalent.
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Protocol, TypeVar
+from typing import Any, Iterable, Mapping, Protocol, Sequence, TypeVar
 from typing_extensions import TypeAlias
 
-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/config/__init__.pyi

@@ -0,0 +1,10 @@
+from typing import NamedTuple
+
+from tensorflow.config import experimental as experimental
+
+class PhysicalDevice(NamedTuple):
+    name: str
+    device_type: str
+
+def list_physical_devices(device_type: None | str = None) -> list[PhysicalDevice]: ...
+def set_visible_devices(devices: list[PhysicalDevice], device_type: None | str = None) -> None: ...

stubs/tensorflow/tensorflow/config/experimental.pyi

@@ -0,0 +1,3 @@
+from tensorflow.config import PhysicalDevice
+
+def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...

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,14 +1,23 @@
 from _typeshed import Incomplete
+from typing import Callable
 
+import tensorflow as tf
+from tensorflow._aliases import _TensorCompatible
 from tensorflow.keras import (
     activations as activations,
     constraints as constraints,
     initializers as initializers,
     layers as layers,
     losses as losses,
     metrics as metrics,
+    models as models,
     optimizers as optimizers,
     regularizers as regularizers,
 )
+from tensorflow.keras.models import Model as Model
 
 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/callbacks.pyi

@@ -0,0 +1,4 @@
+from _typeshed import Incomplete
+
+class Callback(Incomplete): ...
+class History(Incomplete): ...

stubs/tensorflow/tensorflow/keras/layers.pyi

@@ -1,11 +1,11 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Generic, TypeVar, overload
+from typing import Any, Generic, Literal, 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._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: ...
@@ -194,4 +194,53 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]):
         name: str | None = None,
     ) -> None: ...
 
+class Conv2D(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self,
+        filters: int,
+        kernel_size: int | tuple[int, int],
+        strides: int | tuple[int, int] = (1, 1),
+        padding: Literal["valid", "same"] = "valid",
+        data_format: None | Literal["channels_last", "channels_first"] = None,
+        dilation_rate: int | tuple[int, int] = (1, 1),
+        groups: int = 1,
+        activation: _Activation = None,
+        use_bias: bool = True,
+        kernel_initializer: _Initializer = "glorot_uniform",
+        bias_initializer: _Initializer = "zeros",
+        kernel_regularizer: _Regularizer = None,
+        bias_regularizer: _Regularizer = None,
+        activity_regularizer: _Regularizer = None,
+        kernel_constraint: _Constraint = None,
+        bias_constraint: _Constraint = None,
+        trainable: bool = True,
+        dtype: _LayerDtype = None,
+        dynamic: bool = False,
+        name: str | None = None,
+    ) -> None: ...
+
+class Identity(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self, trainable: bool = True, dtype: _LayerDtype = None, dynamic: bool = False, name: str | None = None
+    ) -> None: ...
+
+class LayerNormalization(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self,
+        axis: int = -1,
+        epsilon: float = 0.001,
+        center: bool = True,
+        scale: bool = True,
+        beta_initializer: _Initializer = "zeros",
+        gamma_initializer: _Initializer = "ones",
+        beta_regularizer: _Regularizer = None,
+        gamma_regularizer: _Regularizer = None,
+        beta_constraint: _Constraint = None,
+        gamma_constraint: _Constraint = None,
+        trainable: bool = True,
+        dtype: _LayerDtype = None,
+        dynamic: bool = False,
+        name: str | None = None,
+    ) -> None: ...
+
 def __getattr__(name: str) -> Incomplete: ...

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/metrics.pyi

@@ -1,4 +1,12 @@
+from typing import Any
+from typing_extensions import Self
+
 from tensorflow import Tensor, _TensorCompatible
+from tensorflow.dtypes import DType
+
+class Metric:
+    def __init__(self, name: str, dtype: DType) -> None: ...
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
 
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1