diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 54f489ac7..437a99630 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/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__ diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml index 7922be318..14f4d32a2 100644 --- a/stubs/tensorflow/METADATA.toml +++ b/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"] diff --git a/stubs/tensorflow/tensorflow/data/__init__.pyi b/stubs/tensorflow/tensorflow/data/__init__.pyi index 5dd20bf84..d1b551d95 100644 --- a/stubs/tensorflow/tensorflow/data/__init__.pyi +++ b/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: ... diff --git a/stubs/tensorflow/tensorflow/dtypes.pyi b/stubs/tensorflow/tensorflow/dtypes.pyi index 54bb2846b..a17fb5d4d 100644 --- a/stubs/tensorflow/tensorflow/dtypes.pyi +++ b/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 diff --git a/stubs/tensorflow/tensorflow/feature_column/__init__.pyi b/stubs/tensorflow/tensorflow/feature_column/__init__.pyi index 351cb7658..efae91af0 100644 --- a/stubs/tensorflow/tensorflow/feature_column/__init__.pyi +++ b/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]: ... diff --git a/stubs/tensorflow/tensorflow/keras/optimizers/schedules.pyi b/stubs/tensorflow/tensorflow/keras/optimizers/schedules.pyi index 473314e06..936c7cda1 100644 --- a/stubs/tensorflow/tensorflow/keras/optimizers/schedules.pyi +++ b/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]: ... diff --git a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi b/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi index 08ddc4012..e2f02fa4b 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/feature_column_v2.pyi +++ b/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]: ... diff --git a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi b/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi index 9091f4df6..921ab0865 100644 --- a/stubs/tensorflow/tensorflow/python/feature_column/sequence_feature_column.pyi +++ b/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]: ... diff --git a/stubs/tensorflow/tensorflow/python/framework/dtypes.pyi b/stubs/tensorflow/tensorflow/python/framework/dtypes.pyi new file mode 100644 index 000000000..02577ab52 --- /dev/null +++ b/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 diff --git a/stubs/tensorflow/tensorflow/train/__init__.pyi b/stubs/tensorflow/tensorflow/train/__init__.pyi index 1f42cd14a..591998230 100644 --- a/stubs/tensorflow/tensorflow/train/__init__.pyi +++ b/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: