[tensorflow]: Add a few missing elements (#15265)

Author: hoel-bagard

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,
 )
@@ -37,7 +38,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 (
@@ -441,4 +442,10 @@ 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 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

stubs/tensorflow/tensorflow/keras/metrics.pyi

@@ -1,6 +1,8 @@
+from _typeshed import Incomplete
 from abc import ABCMeta, abstractmethod
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Literal
+from enum import Enum
+from typing import Any, Literal, type_check_only
 from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
@@ -107,6 +109,25 @@ 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"
+    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, sample_weight=None): ...  # type: ignore[override]
+    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