diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index d4cf842f817f..a8961b022b87 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -1,7 +1,7 @@ import abc from _typeshed import Incomplete, Unused from abc import ABC, ABCMeta, abstractmethod -from builtins import bool as _bool +from builtins import bool as _bool, slice as _slice from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from contextlib import contextmanager from enum import Enum @@ -12,8 +12,10 @@ from typing_extensions import Self from google.protobuf.message import Message from tensorflow import ( data as data, + debugging as debugging, experimental as experimental, feature_column as feature_column, + image as image, initializers as initializers, io as io, keras as keras, @@ -26,12 +28,15 @@ from tensorflow._aliases import ( AnyArray, DTypeLike, IntArray, + IntTensorCompatible, RaggedTensorLike, ScalarTensorCompatible, ShapeLike, + Signature, Slice, SparseTensorCompatible, TensorCompatible, + TensorLike, UIntTensorCompatible, ) from tensorflow.autodiff import GradientTape as GradientTape @@ -234,7 +239,7 @@ class TensorShape(metaclass=ABCMeta): @overload def __getitem__(self, key: int) -> int | None: ... @overload - def __getitem__(self, key: slice) -> TensorShape: ... + def __getitem__(self, key: _slice) -> TensorShape: ... def __iter__(self) -> Iterator[int | None]: ... def __len__(self) -> int: ... @@ -375,6 +380,8 @@ class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]): @classmethod def from_value(cls, value: RaggedTensor) -> Self: ... +newaxis: None + def convert_to_tensor( value: TensorCompatible | IndexedSlices, dtype: DTypeLike | None = None, @@ -399,6 +406,7 @@ def squeeze( @overload def squeeze(input: RaggedTensor, axis: int | tuple[int, ...] | list[int], name: str | None = None) -> RaggedTensor: ... +def slice(input_: TensorCompatible, begin: IntTensorCompatible, size: IntTensorCompatible, name: str | None = None) -> Tensor: ... def split( value: TensorCompatible, num_or_size_splits: int | TensorCompatible, @@ -406,6 +414,7 @@ def split( num: int | None = None, name: str | None = "split", ) -> list[Tensor]: ... +def stack(values: TensorCompatible, axis: int = 0, name: str | None = "stack") -> Tensor: ... def tensor_scatter_nd_update( tensor: TensorCompatible, indices: TensorCompatible, updates: TensorCompatible, name: str | None = None ) -> Tensor: ... @@ -442,6 +451,36 @@ def ones_like( ) -> RaggedTensor: ... def reshape(tensor: TensorCompatible, shape: ShapeLike | Tensor, name: str | None = None) -> Tensor: ... +def reverse(tensor: TensorCompatible, axis: IntTensorCompatible, name: str | None = None) -> Tensor: ... + +_ElemT = TypeVar("_ElemT", bound=TensorLike) +_RetT = TypeVar("_RetT", bound=TensorLike) + +@overload +def map_fn( + fn: Callable[[_ElemT], _RetT], + elems: _ElemT, + dtype: DTypeLike | None = None, + parallel_iterations: int | None = None, + back_prop: _bool = True, + swap_memory: _bool = False, + infer_shape: _bool = True, + name: str | None = None, + fn_output_signature: Signature | None = None, +) -> _RetT: ... +@overload +def map_fn( + fn: Callable[[Tensor], _RetT], + elems: TensorCompatible, + dtype: DTypeLike | None = None, + parallel_iterations: int | None = None, + back_prop: _bool = True, + swap_memory: _bool = False, + infer_shape: _bool = True, + name: str | None = None, + fn_output_signature: Signature | None = None, +) -> _RetT: ... + def pad( tensor: TensorCompatible, paddings: Tensor | IntArray | Iterable[Iterable[int]], @@ -467,4 +506,18 @@ def clip_by_value( t: Tensor | IndexedSlices, clip_value_min: TensorCompatible, clip_value_max: TensorCompatible, name: str | None = None ) -> Tensor: ... def tile(input: RaggedTensorLike, multiples: Tensor | Sequence[int], name: str | None = None) -> Tensor: ... + +@overload +def range( + limit: int | Tensor, /, *, delta: int | Tensor = 1, dtype: DTypeLike | None = None, name: str | None = "range" +) -> Tensor: ... +@overload +def range( + start: int | Tensor = 0, + limit: int | Tensor = 0, + delta: int | Tensor = 1, + dtype: DTypeLike | None = None, + name: str | None = "range", +) -> Tensor: ... + def __getattr__(name: str): ... # incomplete module diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index a63f19431717..a3ac8f423437 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -39,6 +39,7 @@ StrDataSequence: TypeAlias = Sequence[str] | Sequence[StrDataSequence] DataSequence: TypeAlias = FloatDataSequence | StrDataSequence | IntDataSequence ScalarTensorCompatible: TypeAlias = tf.Tensor | str | float | np.ndarray[Any, Any] | np.number[Any] UIntTensorCompatible: TypeAlias = tf.Tensor | int | UIntArray +IntTensorCompatible: TypeAlias = tf.Tensor | int | IntArray | Sequence[IntTensorCompatible] FloatTensorCompatible: TypeAlias = tf.Tensor | int | IntArray | float | FloatArray | np.number[Any] StringTensorCompatible: TypeAlias = tf.Tensor | str | npt.NDArray[np.str_] | Sequence[StringTensorCompatible] @@ -56,6 +57,7 @@ TensorOrArray: TypeAlias = tf.Tensor | AnyArray ShapeLike: TypeAlias = tf.TensorShape | Iterable[ScalarTensorCompatible | None] | int | tf.Tensor DTypeLike: TypeAlias = DType | str | np.dtype[Any] | int +Signature: TypeAlias = DType | tf.RaggedTensorSpec | tf.SparseTensorSpec | Sequence[Signature] ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor] ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike] diff --git a/stubs/tensorflow/tensorflow/debugging.pyi b/stubs/tensorflow/tensorflow/debugging.pyi new file mode 100644 index 000000000000..0f48b5e1213d --- /dev/null +++ b/stubs/tensorflow/tensorflow/debugging.pyi @@ -0,0 +1,10 @@ +from tensorflow._aliases import FloatTensorCompatible + +def assert_greater_equal( + x: FloatTensorCompatible, + y: FloatTensorCompatible, + message: str | None = None, + summarize: int | None = None, + name: str | None = None, +) -> None: ... +def __getattr__(name: str): ... # incomplete module diff --git a/stubs/tensorflow/tensorflow/image.pyi b/stubs/tensorflow/tensorflow/image.pyi new file mode 100644 index 000000000000..38fdc40aa059 --- /dev/null +++ b/stubs/tensorflow/tensorflow/image.pyi @@ -0,0 +1,7 @@ +import tensorflow as tf +from tensorflow._aliases import IntTensorCompatible, TensorCompatible + +def random_crop( + value: TensorCompatible, size: IntTensorCompatible, seed: int | None = None, name: str | None = None +) -> tf.Tensor: ... +def __getattr__(name: str): ... # incomplete module