tensorflow: add functions

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,13 +406,15 @@ 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,
     axis: int | Tensor = 0,
     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

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]

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

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