diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 3a356392a6aa..a5d786751db9 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, ) @@ -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 ( @@ -434,4 +435,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 diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi index eba434a93ace..30096af2cafa 100644 --- a/stubs/tensorflow/tensorflow/keras/metrics.pyi +++ b/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