tensorflow: bump version to 2.15 (#11352)

stubs/tensorflow/@tests/stubtest_allowlist.txt

@@ -21,6 +21,12 @@ tensorflow.GradientTape.__getattr__
 tensorflow.data.Dataset.__getattr__
 tensorflow.experimental.Optional.__getattr__
 
+# The Tensor methods below were removed in 2.14, however they are still defined for the
+# internal subclasses that are used at runtime/in practice.
+tensorflow.Tensor.consumers
+tensorflow.Tensor.graph
+tensorflow.Tensor.op
+
 # Internal undocumented API
 tensorflow.RaggedTensor.__init__
 tensorflow.data.Dataset.__init__

stubs/tensorflow/METADATA.toml

@@ -1,4 +1,4 @@
-version = "2.12.*"
+version = "2.15.*"
 upstream_repository = "https://github.com/tensorflow/tensorflow"
 # requires a version of numpy with a `py.typed` file
 requires = ["numpy>=1.20", "types-protobuf"]

stubs/tensorflow/tensorflow/data/__init__.pyi

@@ -1,8 +1,8 @@
 from _typeshed import Incomplete
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterator as _Iterator, Sequence
+from collections.abc import Callable, Collection, Iterator as _Iterator, Sequence
 from typing import Any, Generic, TypeVar, overload
-from typing_extensions import Self
+from typing_extensions import Self, Unpack
 
 import numpy as np
 import tensorflow as tf
@@ -223,8 +223,19 @@ class Dataset(ABC, Generic[_T1]):
         name: str | None = None,
     ) -> Dataset[Dataset[_T1]]: ...
     def with_options(self, options: Options, name: str | None = None) -> Dataset[_T1]: ...
+    @overload
     @staticmethod
-    def zip(datasets: tuple[Dataset[_T2], Dataset[_T3]], name: str | None = None) -> Dataset[tuple[_T2, _T3]]: ...
+    def zip(
+        *args: Collection[Dataset[Any]] | ContainerGeneric[Dataset[Any]], name: str | None = None
+    ) -> Dataset[tuple[Any, ...]]: ...
+    @overload
+    @staticmethod
+    def zip(*args: Unpack[tuple[Dataset[_T2], Dataset[_T3]]], name: str | None = None) -> Dataset[tuple[_T2, _T3]]: ...
+    @overload
+    @staticmethod
+    def zip(
+        *, datasets: tuple[Dataset[_T2], Dataset[_T3]] | None = None, name: str | None = None
+    ) -> Dataset[tuple[_T2, _T3]]: ...
     def __len__(self) -> int: ...
     def __nonzero__(self) -> bool: ...
     def __getattr__(self, name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/dtypes.pyi

@@ -5,10 +5,12 @@ from typing import Any
 
 import numpy as np
 from tensorflow._aliases import DTypeLike
+from tensorflow.python.framework.dtypes import HandleData
 
 class _DTypeMeta(ABCMeta): ...
 
 class DType(metaclass=_DTypeMeta):
+    def __init__(self, type_enum: int, handle_data: HandleData | None = None) -> None: ...
     @property
     def name(self) -> str: ...
     @property

stubs/tensorflow/tensorflow/feature_column/__init__.pyi

@@ -91,5 +91,5 @@ def sequence_categorical_column_with_vocabulary_list(
     num_oov_buckets: int = 0,
 ) -> fc.SequenceCategoricalColumn: ...
 def make_parse_example_spec(
-    feature_columns: Iterable[fc.FeatureColumn],
+    feature_columns: Iterable[fc._FeatureColumn],
 ) -> dict[str, tf.io.FixedLenFeature | tf.io.VarLenFeature]: ...

stubs/tensorflow/tensorflow/keras/optimizers/schedules.pyi

@@ -57,7 +57,13 @@ class PolynomialDecay(LearningRateSchedule):
 
 class CosineDecay(LearningRateSchedule):
     def __init__(
-        self, initial_learning_rate: float | tf.Tensor, decay_steps: int, alpha: float | tf.Tensor = 0.0, name: str | None = None
+        self,
+        initial_learning_rate: float | tf.Tensor,
+        decay_steps: int,
+        alpha: float | tf.Tensor = 0.0,
+        name: str | None = None,
+        warmup_target: int | tf.Tensor | None = None,  # float32 or float64 Tensor
+        warmup_steps: int | tf.Tensor = 0,  # int32 or int64 Tensor
     ) -> None: ...
     def __call__(self, step: int | tf.Tensor) -> float | tf.Tensor: ...
     def get_config(self) -> dict[str, Any]: ...

stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi

@@ -14,21 +14,21 @@ from tensorflow._aliases import ShapeLike
 _Combiners: TypeAlias = Literal["mean", "sqrtn", "sum"]
 _ExampleSpec: TypeAlias = dict[str, tf.io.FixedLenFeature | tf.io.VarLenFeature]
 
-class FeatureColumn(ABC):
+class _FeatureColumn(ABC):
     @property
     @abstractmethod
     def name(self) -> str: ...
     @property
     @abstractmethod
     def parse_example_spec(self) -> _ExampleSpec: ...
-    def __lt__(self, other: FeatureColumn) -> bool: ...
-    def __gt__(self, other: FeatureColumn) -> bool: ...
+    def __lt__(self, other: _FeatureColumn) -> bool: ...
+    def __gt__(self, other: _FeatureColumn) -> bool: ...
     @property
     @abstractmethod
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
-class DenseColumn(FeatureColumn, metaclass=ABCMeta): ...
-class SequenceDenseColumn(FeatureColumn, metaclass=ABCMeta): ...
+class DenseColumn(_FeatureColumn, metaclass=ABCMeta): ...
+class SequenceDenseColumn(_FeatureColumn, metaclass=ABCMeta): ...
 
 # These classes are mostly subclasses of collections.namedtuple but we can't use
 # typing.NamedTuple because they use multiple inheritance with other non namedtuple classes.
@@ -53,9 +53,9 @@ class NumericColumn(DenseColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
-class CategoricalColumn(FeatureColumn):
+class CategoricalColumn(_FeatureColumn):
     @property
     @abstractmethod
     def num_buckets(self) -> int: ...
@@ -72,7 +72,7 @@ class BucketizedColumn(DenseColumn, CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class EmbeddingColumn(DenseColumn, SequenceDenseColumn):
     categorical_column: CategoricalColumn
@@ -103,7 +103,7 @@ class EmbeddingColumn(DenseColumn, SequenceDenseColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class SharedEmbeddingColumnCreator:
     def __init__(
@@ -139,7 +139,7 @@ class SharedEmbeddingColumn(DenseColumn, SequenceDenseColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class CrossedColumn(CategoricalColumn):
     keys: tuple[str, ...]
@@ -154,7 +154,7 @@ class CrossedColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class IdentityCategoricalColumn(CategoricalColumn):
     key: str
@@ -169,7 +169,7 @@ class IdentityCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class HashedCategoricalColumn(CategoricalColumn):
     key: str
@@ -184,7 +184,7 @@ class HashedCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class VocabularyFileCategoricalColumn(CategoricalColumn):
     key: str
@@ -212,7 +212,7 @@ class VocabularyFileCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class VocabularyListCategoricalColumn(CategoricalColumn):
     key: str
@@ -231,7 +231,7 @@ class VocabularyListCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class WeightedCategoricalColumn(CategoricalColumn):
     categorical_column: CategoricalColumn
@@ -246,7 +246,7 @@ class WeightedCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class IndicatorColumn(DenseColumn, SequenceDenseColumn):
     categorical_column: CategoricalColumn
@@ -257,7 +257,7 @@ class IndicatorColumn(DenseColumn, SequenceDenseColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...
 
 class SequenceCategoricalColumn(CategoricalColumn):
     categorical_column: CategoricalColumn
@@ -270,4 +270,4 @@ class SequenceCategoricalColumn(CategoricalColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...

stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi

@@ -3,7 +3,7 @@ from typing_extensions import Self
 
 import tensorflow as tf
 from tensorflow._aliases import ShapeLike
-from tensorflow.python.feature_column.feature_column_v2 import FeatureColumn, SequenceDenseColumn, _ExampleSpec
+from tensorflow.python.feature_column.feature_column_v2 import SequenceDenseColumn, _ExampleSpec, _FeatureColumn
 
 # Strangely at runtime most of Sequence feature columns are defined in feature_column_v2 except
 # for this one.
@@ -27,4 +27,4 @@ class SequenceNumericColumn(SequenceDenseColumn):
     @property
     def parse_example_spec(self) -> _ExampleSpec: ...
     @property
-    def parents(self) -> list[FeatureColumn | str]: ...
+    def parents(self) -> list[_FeatureColumn | str]: ...

stubs/tensorflow/tensorflow/python/framework/dtypes.pyi

@@ -0,0 +1,7 @@
+import dataclasses
+from _typeshed import Incomplete
+
+@dataclasses.dataclass(frozen=True)
+class HandleData:
+    shape_inference: Incomplete | None = None
+    alias_id: int | None = None

stubs/tensorflow/tensorflow/train/__init__.pyi

@@ -20,16 +20,14 @@ from tensorflow.python.trackable.base import Trackable
 class CheckpointOptions:
     experimental_io_device: None | str
     experimental_enable_async_checkpoint: bool
-    # Uncomment when the stubs' TF version is updated to 2.15
-    # experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]]
+    experimental_write_callbacks: None | list[Callable[[str], object] | Callable[[], object]]
     enable_async: bool
 
     def __init__(
         self,
         experimental_io_device: None | str = None,
         experimental_enable_async_checkpoint: bool = False,
-        # Uncomment when the stubs' TF version is updated to 2.15
-        # experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None,
+        experimental_write_callbacks: None | list[Callable[[str], object] | Callable[[], object]] = None,
         enable_async: bool = False,
     ) -> None: ...
 
@@ -51,7 +49,7 @@ class Checkpoint:
     def read(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...
     def restore(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...
     def save(self, file_prefix: str, options: CheckpointOptions | None = None) -> str: ...
-    # def sync(self) -> None: ...  # Uncomment when the stubs' TF version is updated to 2.15
+    def sync(self) -> None: ...
     def write(self, file_prefix: str, options: CheckpointOptions | None = None) -> str: ...
 
 class CheckpointManager: