From f80f9c293ce95486bdcc1c5829903f91dde01af0 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 28 Dec 2024 22:15:54 +0900 Subject: [PATCH 1/7] add tensorflow.math.round --- stubs/tensorflow/tensorflow/__init__.pyi | 1 + stubs/tensorflow/tensorflow/math.pyi | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 21c8df673cbd..7f29f1d79f38 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -56,6 +56,7 @@ from tensorflow.math import ( reduce_min as reduce_min, reduce_prod as reduce_prod, reduce_sum as reduce_sum, + round as round, sigmoid as sigmoid, sign as sign, sin as sin, diff --git a/stubs/tensorflow/tensorflow/math.pyi b/stubs/tensorflow/tensorflow/math.pyi index d3b02bf7a748..2e72ada0a8f2 100644 --- a/stubs/tensorflow/tensorflow/math.pyi +++ b/stubs/tensorflow/tensorflow/math.pyi @@ -219,6 +219,12 @@ def square(x: RaggedTensor, name: str | None = None) -> RaggedTensor: ... def softplus(features: TensorCompatible, name: str | None = None) -> Tensor: ... @overload def softplus(features: RaggedTensor, name: str | None = None) -> RaggedTensor: ... +@overload +def round(x: TensorCompatible, name: str | None = None) -> Tensor: ... +@overload +def round(x: SparseTensor, name: str | None = None) -> SparseTensor: ... +@overload +def round(x: RaggedTensor, name: str | None = None) -> RaggedTensor: ... # Depending on the method axis is either a rank 0 tensor or a rank 0/1 tensor. def reduce_mean( From 06f5932136e7881944b381d8353b746a1faa3dca Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 28 Dec 2024 22:28:30 +0900 Subject: [PATCH 2/7] add tensorflow.pad fix typo --- stubs/tensorflow/tensorflow/__init__.pyi | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 7f29f1d79f38..93959496b317 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -6,7 +6,7 @@ from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from contextlib import contextmanager from enum import Enum from types import TracebackType -from typing import Any, Generic, NoReturn, TypeVar, overload +from typing import Any, Generic, Literal, NoReturn, TypeVar, overload from typing_extensions import ParamSpec, Self from google.protobuf.message import Message @@ -20,7 +20,7 @@ from tensorflow import ( math as math, types as types, ) -from tensorflow._aliases import AnyArray, DTypeLike, ShapeLike, Slice, TensorCompatible +from tensorflow._aliases import AnyArray, DTypeLike, ScalarTensorCompatible, ShapeLike, Slice, TensorCompatible from tensorflow.autodiff import GradientTape as GradientTape from tensorflow.core.protobuf import struct_pb2 from tensorflow.dtypes import * @@ -404,4 +404,11 @@ def ones_like( input: RaggedTensor, dtype: DTypeLike | None = None, name: str | None = None, layout: Layout | None = None ) -> RaggedTensor: ... def reshape(tensor: TensorCompatible, shape: ShapeLike | Tensor, name: str | None = None) -> Tensor: ... +def pad( + tensor: TensorCompatible, + paddings: ShapeLike, + mode: Literal["CONSTANT", "constant", "REFLECT", "reflect", "SYMMETRIC", "symmectric"] = "CONSTANT", + constant_values: ScalarTensorCompatible = 0, + name: str | None = None, +) -> Tensor: ... def __getattr__(name: str) -> Incomplete: ... From 89d1bcb2a595287d2d0735199b08165efaee76b1 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 29 Dec 2024 10:11:15 +0900 Subject: [PATCH 3/7] add tf.shape --- 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 93959496b317..905a40455da3 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -411,4 +411,5 @@ def pad( constant_values: ScalarTensorCompatible = 0, name: str | None = None, ) -> Tensor: ... +def shape(input: TensorCompatible, out_type: DTypeLike | None = None, name: str | None = None) -> Tensor: ... def __getattr__(name: str) -> Incomplete: ... From 4e678f2df12ff984cafae10527fb3b8df555238d Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 29 Dec 2024 10:14:30 +0900 Subject: [PATCH 4/7] add tf.where --- stubs/tensorflow/tensorflow/__init__.pyi | 3 +++ 1 file changed, 3 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 905a40455da3..bc95b3512830 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -412,4 +412,7 @@ def pad( name: str | None = None, ) -> Tensor: ... def shape(input: TensorCompatible, out_type: DTypeLike | None = None, name: str | None = None) -> Tensor: ... +def where( + condition: TensorCompatible, x: TensorCompatible | None = None, y: TensorCompatible | None = None, name: str | None = None +) -> Tensor: ... def __getattr__(name: str) -> Incomplete: ... From 3dd5cbfc68d2dff320c2bacd15f18a21cdba7ab3 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 29 Dec 2024 10:19:23 +0900 Subject: [PATCH 5/7] add tensorflow.gather_nd add bad_indices_policy arg to gather_nd --- stubs/tensorflow/tensorflow/__init__.pyi | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index bc95b3512830..64b2c249b61d 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -20,7 +20,15 @@ from tensorflow import ( math as math, types as types, ) -from tensorflow._aliases import AnyArray, DTypeLike, ScalarTensorCompatible, ShapeLike, Slice, TensorCompatible +from tensorflow._aliases import ( + AnyArray, + DTypeLike, + ScalarTensorCompatible, + ShapeLike, + Slice, + TensorCompatible, + UIntTensorCompatible, +) from tensorflow.autodiff import GradientTape as GradientTape from tensorflow.core.protobuf import struct_pb2 from tensorflow.dtypes import * @@ -415,4 +423,11 @@ def shape(input: TensorCompatible, out_type: DTypeLike | None = None, name: str def where( condition: TensorCompatible, x: TensorCompatible | None = None, y: TensorCompatible | None = None, name: str | None = None ) -> Tensor: ... +def gather_nd( + params: TensorCompatible, + indices: UIntTensorCompatible, + batch_dims: UIntTensorCompatible = 0, + name: str | None = None, + bad_indices_policy: Literal["", "DEFAULT", "ERROR", "IGNORE"] = "", +) -> Tensor: ... def __getattr__(name: str) -> Incomplete: ... From 43304b7dd0ec817e0364e3a9789021cb009f3fbc Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 29 Dec 2024 10:27:16 +0900 Subject: [PATCH 6/7] add tensorflow.signal.hamming_window --- stubs/tensorflow/tensorflow/signal.pyi | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 stubs/tensorflow/tensorflow/signal.pyi diff --git a/stubs/tensorflow/tensorflow/signal.pyi b/stubs/tensorflow/tensorflow/signal.pyi new file mode 100644 index 000000000000..c0ce254f1e36 --- /dev/null +++ b/stubs/tensorflow/tensorflow/signal.pyi @@ -0,0 +1,6 @@ +from tensorflow import Tensor +from tensorflow._aliases import DTypeLike, TensorCompatible + +def hamming_window( + window_length: TensorCompatible, periodic: bool | TensorCompatible = True, dtype: DTypeLike = ..., name: str | None = None +) -> Tensor: ... From 5f233ecb61f6c45ae6302af12439d6ec0f8f1161 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 5 Jan 2025 13:07:48 +0900 Subject: [PATCH 7/7] fix paddings' type --- stubs/tensorflow/tensorflow/__init__.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 64b2c249b61d..b571ab0d82f6 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -23,6 +23,7 @@ from tensorflow import ( from tensorflow._aliases import ( AnyArray, DTypeLike, + IntArray, ScalarTensorCompatible, ShapeLike, Slice, @@ -414,7 +415,7 @@ def ones_like( def reshape(tensor: TensorCompatible, shape: ShapeLike | Tensor, name: str | None = None) -> Tensor: ... def pad( tensor: TensorCompatible, - paddings: ShapeLike, + paddings: Tensor | IntArray | Iterable[Iterable[int]], mode: Literal["CONSTANT", "constant", "REFLECT", "reflect", "SYMMETRIC", "symmectric"] = "CONSTANT", constant_values: ScalarTensorCompatible = 0, name: str | None = None,