From e29b7c4e3e28e8c13c2293ad7aef58cf67477978 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 27 Jul 2026 09:02:23 +0900 Subject: [PATCH 01/11] add tf.range --- stubs/tensorflow/tensorflow/__init__.pyi | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index d4cf842f817f..93be0b4485cd 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -467,4 +467,16 @@ 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 From eb5e601edba7431fe102f2cad00042f6f4b40037 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 27 Jul 2026 09:15:22 +0900 Subject: [PATCH 02/11] add tf.debugging.assert_greater_equal --- stubs/tensorflow/tensorflow/debugging.pyi | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 stubs/tensorflow/tensorflow/debugging.pyi diff --git a/stubs/tensorflow/tensorflow/debugging.pyi b/stubs/tensorflow/tensorflow/debugging.pyi new file mode 100644 index 000000000000..981f80faff29 --- /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, +): ... +def __getattr__(name: str): ... # incomplete module From 5dca5bdbcd76163d1e0ed6f37fb78c565b0258d4 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 27 Jul 2026 10:37:00 +0900 Subject: [PATCH 03/11] add tf.image.random_crop --- stubs/tensorflow/tensorflow/image.pyi | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 stubs/tensorflow/tensorflow/image.pyi diff --git a/stubs/tensorflow/tensorflow/image.pyi b/stubs/tensorflow/tensorflow/image.pyi new file mode 100644 index 000000000000..71409c00fd1a --- /dev/null +++ b/stubs/tensorflow/tensorflow/image.pyi @@ -0,0 +1,9 @@ +from collections.abc import Sequence + +from tensorflow import Tensor +from tensorflow._aliases import IntArray, TensorCompatible + +def random_crop( + value: TensorCompatible, size: Sequence[int] | IntArray | Tensor, seed: int | None = None, name: str | None = None +): ... +def __getattr__(name: str): ... # incomplete module From 4801fc345f1908ea1c5d7ae13a9c5815c1169503 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 27 Jul 2026 10:52:17 +0900 Subject: [PATCH 04/11] add tf.map_fn --- stubs/tensorflow/tensorflow/__init__.pyi | 12 ++++++++++++ stubs/tensorflow/tensorflow/_aliases.pyi | 1 + 2 files changed, 13 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 93be0b4485cd..c924414f7f7c 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -29,6 +29,7 @@ from tensorflow._aliases import ( RaggedTensorLike, ScalarTensorCompatible, ShapeLike, + Signature, Slice, SparseTensorCompatible, TensorCompatible, @@ -442,6 +443,17 @@ def ones_like( ) -> RaggedTensor: ... def reshape(tensor: TensorCompatible, shape: ShapeLike | Tensor, name: str | None = None) -> Tensor: ... +def map_fn( + fn: Callable[TensorLike | TensorCompatible, TensorLike], + elems: TensorLike | 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, +) -> Tensor: ... def pad( tensor: TensorCompatible, paddings: Tensor | IntArray | Iterable[Iterable[int]], diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index a63f19431717..0d3a37dd1844 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -56,6 +56,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: DType | tf.RaggedTensorSpec | tf.SparseTensorSpec | Sequence[Signature] ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor] ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike] From 41067f029b05088804ded5ef17dffdf2558f113e Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 27 Jul 2026 10:59:12 +0900 Subject: [PATCH 05/11] add tf.newaxis --- stubs/tensorflow/tensorflow/__init__.pyi | 2 ++ 1 file changed, 2 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index c924414f7f7c..377c19967918 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -376,6 +376,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, From a971b2a64feaeb969fa17448a4fcb613c8b550e8 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 07:55:55 +0900 Subject: [PATCH 06/11] add tf.slice --- stubs/tensorflow/tensorflow/__init__.pyi | 1 + stubs/tensorflow/tensorflow/_aliases.pyi | 1 + 2 files changed, 2 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 377c19967918..5e30247ba82d 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -402,6 +402,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, diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index 0d3a37dd1844..81b421b4314d 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 FloatTensorCompatible: TypeAlias = tf.Tensor | int | IntArray | float | FloatArray | np.number[Any] StringTensorCompatible: TypeAlias = tf.Tensor | str | npt.NDArray[np.str_] | Sequence[StringTensorCompatible] From 410287e382f87f60efde853812f6058f1a61093a Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 08:00:27 +0900 Subject: [PATCH 07/11] add tf.stack --- stubs/tensorflow/tensorflow/__init__.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 5e30247ba82d..0f4b17d46162 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -410,6 +410,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: ... From e1a09bfbe403fcbbb3b212edb9d0539a886fc79f Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 08:05:33 +0900 Subject: [PATCH 08/11] add tf.reverse --- stubs/tensorflow/tensorflow/__init__.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 0f4b17d46162..6a8275905334 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -447,6 +447,7 @@ 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: ... def map_fn( fn: Callable[TensorLike | TensorCompatible, TensorLike], elems: TensorLike | TensorCompatible, From 9c02549662662a2c57843cd5d25eb75459a4a7cd Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 08:14:42 +0900 Subject: [PATCH 09/11] fix: IntTensorCompatible type --- stubs/tensorflow/tensorflow/_aliases.pyi | 2 +- stubs/tensorflow/tensorflow/image.pyi | 9 ++------- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index 81b421b4314d..14d4dbc12191 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -39,7 +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 +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] diff --git a/stubs/tensorflow/tensorflow/image.pyi b/stubs/tensorflow/tensorflow/image.pyi index 71409c00fd1a..e99605f5da78 100644 --- a/stubs/tensorflow/tensorflow/image.pyi +++ b/stubs/tensorflow/tensorflow/image.pyi @@ -1,9 +1,4 @@ -from collections.abc import Sequence +from tensorflow._aliases import IntTensorCompatible, TensorCompatible -from tensorflow import Tensor -from tensorflow._aliases import IntArray, TensorCompatible - -def random_crop( - value: TensorCompatible, size: Sequence[int] | IntArray | Tensor, seed: int | None = None, name: str | None = None -): ... +def random_crop(value: TensorCompatible, size: IntTensorCompatible, seed: int | None = None, name: str | None = None): ... def __getattr__(name: str): ... # incomplete module From f6ccdabeec384b12567c856a778091c99e8fb69c Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 09:06:23 +0900 Subject: [PATCH 10/11] fix: add missing imports, re-exports, return values Also a few type fixes --- stubs/tensorflow/tensorflow/__init__.pyi | 36 ++++++++++++++++++----- stubs/tensorflow/tensorflow/_aliases.pyi | 2 +- stubs/tensorflow/tensorflow/debugging.pyi | 2 +- stubs/tensorflow/tensorflow/image.pyi | 5 +++- 4 files changed, 35 insertions(+), 10 deletions(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 6a8275905334..b92d4ab8e6f8 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -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,6 +28,7 @@ from tensorflow._aliases import ( AnyArray, DTypeLike, IntArray, + IntTensorCompatible, RaggedTensorLike, ScalarTensorCompatible, ShapeLike, @@ -33,6 +36,7 @@ from tensorflow._aliases import ( Slice, SparseTensorCompatible, TensorCompatible, + TensorLike, UIntTensorCompatible, ) from tensorflow.autodiff import GradientTape as GradientTape @@ -448,17 +452,35 @@ def ones_like( 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[TensorLike | TensorCompatible, TensorLike], - elems: TensorLike | TensorCompatible, + 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, + back_prop: _bool = True, + swap_memory: _bool = False, + infer_shape: _bool = True, name: str | None = None, - fn_output_signature: Signature = None, -) -> Tensor: ... + 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]], diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index 14d4dbc12191..a3ac8f423437 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -57,7 +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: DType | tf.RaggedTensorSpec | tf.SparseTensorSpec | Sequence[Signature] +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 index 981f80faff29..0f48b5e1213d 100644 --- a/stubs/tensorflow/tensorflow/debugging.pyi +++ b/stubs/tensorflow/tensorflow/debugging.pyi @@ -6,5 +6,5 @@ def assert_greater_equal( 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 index e99605f5da78..38fdc40aa059 100644 --- a/stubs/tensorflow/tensorflow/image.pyi +++ b/stubs/tensorflow/tensorflow/image.pyi @@ -1,4 +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): ... +def random_crop( + value: TensorCompatible, size: IntTensorCompatible, seed: int | None = None, name: str | None = None +) -> tf.Tensor: ... def __getattr__(name: str): ... # incomplete module From 69a143de1fceb16b09c848d1c52de041f560275c Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Tue, 28 Jul 2026 12:15:28 +0900 Subject: [PATCH 11/11] fix: builtin shadowing --- stubs/tensorflow/tensorflow/__init__.pyi | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index b92d4ab8e6f8..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 @@ -239,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: ... @@ -506,6 +506,7 @@ 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" @@ -518,4 +519,5 @@ def range( dtype: DTypeLike | None = None, name: str | None = "range", ) -> Tensor: ... + def __getattr__(name: str): ... # incomplete module