From 0cf2521d6ba1eeea0f3678657b754230b1cc0659 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 12:48:59 +0900 Subject: [PATCH 01/10] tensorflow: add missing matmul re-export --- stubs/tensorflow/tensorflow/__init__.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 3a356392a6aa..889aac1550f4 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -37,7 +37,7 @@ from tensorflow.core.protobuf import struct_pb2 from tensorflow.dtypes import * from tensorflow.experimental.dtensor import Layout from tensorflow.keras import losses as losses -from tensorflow.linalg import eye as eye +from tensorflow.linalg import eye as eye, matmul as matmul # Most tf.math functions are exported as tf, but sadly not all are. from tensorflow.math import ( From 36598674cf73acd6f849c112afdefd5c2cedcda0 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 21:38:48 +0900 Subject: [PATCH 02/10] tensorflow: add tf.transpose --- 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 889aac1550f4..6953c29a2fe9 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -434,4 +434,7 @@ def gather_nd( name: str | None = None, bad_indices_policy: Literal["", "DEFAULT", "ERROR", "IGNORE"] = "", ) -> Tensor: ... +def transpose( + a: Tensor, perm: Sequence[int] | IntArray | None = None, conjugate: _bool = False, name: str = "transpose" +) -> Tensor: ... def __getattr__(name: str): ... # incomplete module From 1585d031aa63b36b1da3351008b5f9985cbe2551 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 21:44:53 +0900 Subject: [PATCH 03/10] tensorflow: add missing tf.nn re-export --- 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 6953c29a2fe9..2efdfc02aeb2 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -18,6 +18,7 @@ from tensorflow import ( io as io, keras as keras, math as math, + nn as nn, random as random, types as types, ) From 85474e7a92ac091a1377530190c028dbf781e108 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 21:59:22 +0900 Subject: [PATCH 04/10] tensorflow: add tf.clip_by_value --- 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 2efdfc02aeb2..a5d786751db9 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -438,4 +438,7 @@ def gather_nd( def transpose( a: Tensor, perm: Sequence[int] | IntArray | None = None, conjugate: _bool = False, name: str = "transpose" ) -> Tensor: ... +def clip_by_value( + t: Tensor | IndexedSlices, clip_value_min: TensorCompatible, clip_value_max: TensorCompatible, name: str | None = None +) -> Tensor: ... def __getattr__(name: str): ... # incomplete module From 0ec1986420da41b0232605d274846d1eefc33d87 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 22:04:10 +0900 Subject: [PATCH 05/10] tensorflow: add tf.keras.metrics.Mean --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 3 +++ 1 file changed, 3 insertions(+) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index eba434a93ace..e2d56ad7c63c 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -99,6 +99,9 @@ class Accuracy(MeanMetricWrapper): class CategoricalAccuracy(MeanMetricWrapper): def __init__(self, name: str | None = "categorical_accuracy", dtype: DTypeLike | None = None) -> None: ... +class Mean(MeanMetricWrapper): + def __init__(self, name: str | None = "mean", dtype: DTypeLike | None = None) -> None: ... + class TopKCategoricalAccuracy(MeanMetricWrapper): def __init__(self, k: int = 5, name: str | None = "top_k_categorical_accuracy", dtype: DTypeLike | None = None) -> None: ... From bfc1bdb5321afc49564c8e20dd7939f2be04cf60 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 22:47:55 +0900 Subject: [PATCH 06/10] fix: Mean base class --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index e2d56ad7c63c..b8cbf59d9629 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod from collections.abc import Callable, Iterable, Sequence +from enum import Enum from typing import Any, Literal from typing_extensions import Self, TypeAlias @@ -99,9 +100,6 @@ class Accuracy(MeanMetricWrapper): class CategoricalAccuracy(MeanMetricWrapper): def __init__(self, name: str | None = "categorical_accuracy", dtype: DTypeLike | None = None) -> None: ... -class Mean(MeanMetricWrapper): - def __init__(self, name: str | None = "mean", dtype: DTypeLike | None = None) -> None: ... - class TopKCategoricalAccuracy(MeanMetricWrapper): def __init__(self, k: int = 5, name: str | None = "top_k_categorical_accuracy", dtype: DTypeLike | None = None) -> None: ... @@ -110,6 +108,19 @@ class SparseTopKCategoricalAccuracy(MeanMetricWrapper): self, k: int = 5, name: str | None = "sparse_top_k_categorical_accuracy", dtype: DTypeLike | None = None ) -> None: ... +class _Reduction(Enum): + SUM = "sum" + SUM_OVER_BATCH_SIZE = "sum_over_batch_size" + WEIGHTED_MEAN = "weighted_mean" + +class Reduce(Metric): + def __init__(self, reduction: _Reduction, name: str | None, dtype: DTypeLike | None = None) -> None: ... + def update_state(self, values: TensorCompatible, sample_weight: TensorCompatible | None = None) -> Operation: ... + def result(self) -> Tensor: ... + +class Mean(Reduce): + def __init__(self, name: str | None = "mean", dtype: DTypeLike | None = None) -> None: ... + def serialize(metric: KerasSerializable) -> dict[str, Any]: ... def binary_crossentropy( y_true: TensorCompatible, y_pred: TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1 From ce772a4afb6b0539de4f59da6decb498216f5624 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ho=C3=ABl=20Bagard?= <34478245+hoel-bagard@users.noreply.github.com> Date: Mon, 12 Jan 2026 22:59:58 +0900 Subject: [PATCH 07/10] Update stubs/tensorflow/tensorflow/keras/metrics.pyi Co-authored-by: Sebastian Rittau --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 3 +++ 1 file changed, 3 insertions(+) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index b8cbf59d9629..44c4b529a7b5 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -108,6 +108,9 @@ class SparseTopKCategoricalAccuracy(MeanMetricWrapper): self, k: int = 5, name: str | None = "sparse_top_k_categorical_accuracy", dtype: DTypeLike | None = None ) -> None: ... +# TODO: Actually tensorflow.python.keras.utils.metrics_utils.Reduction, but that module +# is currently missing from the stub. +@type_check_only class _Reduction(Enum): SUM = "sum" SUM_OVER_BATCH_SIZE = "sum_over_batch_size" From 1ab5f788cfe0993627192d42d82e02ebb75403da Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 23:01:33 +0900 Subject: [PATCH 08/10] fix: add Reduce class attributes --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index 44c4b529a7b5..90e5475d2c92 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -1,3 +1,4 @@ +from _typeshed import Incomplete from abc import ABCMeta, abstractmethod from collections.abc import Callable, Iterable, Sequence from enum import Enum @@ -117,6 +118,9 @@ class _Reduction(Enum): WEIGHTED_MEAN = "weighted_mean" class Reduce(Metric): + reduction: _Reduction + total: Incomplete + count: Incomplete # only defined for some reductions def __init__(self, reduction: _Reduction, name: str | None, dtype: DTypeLike | None = None) -> None: ... def update_state(self, values: TensorCompatible, sample_weight: TensorCompatible | None = None) -> Operation: ... def result(self) -> Tensor: ... From 22957ffc1398442adba1dd29441ba260eb88c412 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ho=C3=ABl=20Bagard?= <34478245+hoel-bagard@users.noreply.github.com> Date: Mon, 12 Jan 2026 23:03:38 +0900 Subject: [PATCH 09/10] add missing type_check_only import Co-authored-by: Sebastian Rittau --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index 90e5475d2c92..964fce86b94c 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -2,7 +2,7 @@ from _typeshed import Incomplete from abc import ABCMeta, abstractmethod from collections.abc import Callable, Iterable, Sequence from enum import Enum -from typing import Any, Literal +from typing import Any, Literal, type_check_only from typing_extensions import Self, TypeAlias import tensorflow as tf From 93109a3cd28839489cf83129428c064ad27c5d2c Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Mon, 12 Jan 2026 23:08:13 +0900 Subject: [PATCH 10/10] fix ci --- stubs/tensorflow/tensorflow/keras/metrics.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index 964fce86b94c..30096af2cafa 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi @@ -122,7 +122,7 @@ class Reduce(Metric): total: Incomplete count: Incomplete # only defined for some reductions def __init__(self, reduction: _Reduction, name: str | None, dtype: DTypeLike | None = None) -> None: ... - def update_state(self, values: TensorCompatible, sample_weight: TensorCompatible | None = None) -> Operation: ... + def update_state(self, values, sample_weight=None): ... # type: ignore[override] def result(self) -> Tensor: ... class Mean(Reduce):