From ba67d5ba5158751d6231fe5cfcabfbabc43afb47 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Fri, 19 Jan 2024 16:57:40 +0900
Subject: [PATCH 01/77] Add convert_to_tensor

---
 stubs/tensorflow/tensorflow/__init__.pyi | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 18f49fe372ae..400b2b77a39e 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -406,9 +406,10 @@ class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
     def from_value(cls, value: RaggedTensor) -> Self: ...
 
 def __getattr__(name: str) -> Incomplete: ...
+
 def convert_to_tensor(
-    value: _TensorCompatible | IndexedSlices,
-    dtype: _DTypeLike | None = None,
-    dtype_hint: _DTypeLike | None = None,
-    name: str | None = None,
+    value: _TensorCompatible,
+    dtype: DType | None = None,
+    dtype_hint: DType | None = None,
+    name: str | None = None
 ) -> Tensor: ...

From c9940633666629d81e8ecc67e2969402b909ea64 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 19 Jan 2024 08:13:38 +0000
Subject: [PATCH 02/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/__init__.pyi | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 400b2b77a39e..8bff6e3b4734 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -406,10 +406,6 @@ class RaggedTensorSpec(TypeSpec[struct_pb2.TypeSpecProto]):
     def from_value(cls, value: RaggedTensor) -> Self: ...
 
 def __getattr__(name: str) -> Incomplete: ...
-
 def convert_to_tensor(
-    value: _TensorCompatible,
-    dtype: DType | None = None,
-    dtype_hint: DType | None = None,
-    name: str | None = None
+    value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...

From 476c5f6da592a746edb74a56cb77b867d0f24be4 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Mon, 22 Jan 2024 08:18:25 +0900
Subject: [PATCH 03/77] Add a few functions.

---
 stubs/tensorflow/tensorflow/__init__.pyi       | 4 ++++
 stubs/tensorflow/tensorflow/keras/__init__.pyi | 1 +
 stubs/tensorflow/tensorflow/keras/models.pyi   | 3 +++
 3 files changed, 8 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/keras/models.pyi

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 8bff6e3b4734..5b0e66671c25 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -409,3 +409,7 @@ def __getattr__(name: str) -> Incomplete: ...
 def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
+def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
+# TODO: The "list"s below are from the TF doc, any Iterable or Sequence might be enough ?
+def concat(values: list[Tensor] | Tensor, axis: int, name: None | str = "concat") -> Tensor: ...
+def squeeze(input: Tensor, axis: None | list[int] = None, name: None | str = None) -> Tensor: ...
diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 09b1b4737528..5ea092c24306 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -9,6 +9,7 @@ from tensorflow.keras import (
     metrics as metrics,
     optimizers as optimizers,
     regularizers as regularizers,
+    models as models,
 )
 
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
new file mode 100644
index 000000000000..787563f5c769
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -0,0 +1,3 @@
+# class Model(base_layer.Layer, version_utils.ModelVersionSelector):
+class Model():
+    def __init__(self, *args, **kwargs) -> None: ...

From 8a8df5418433f6f4d5a2beba9c598aff0e553416 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Mon, 22 Jan 2024 10:34:43 +0900
Subject: [PATCH 04/77] feat: add config functions.

---
 stubs/tensorflow/tensorflow/config.pyi | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/config.pyi

diff --git a/stubs/tensorflow/tensorflow/config.pyi b/stubs/tensorflow/tensorflow/config.pyi
new file mode 100644
index 000000000000..14f8099bd2b4
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/config.pyi
@@ -0,0 +1,9 @@
+from typing import NamedTuple
+
+class PhysicalDevice(NamedTuple):
+    name: str
+    device_type: str
+
+def list_physical_devices(device_type: None | str = None) -> list[PhysicalDevice]: ...
+def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...
+def set_visible_devices(devices: list[PhysicalDevice], device_type: None | str = None) -> None: ...

From 9f69588e738cec080a58d8285e98d99fb04f35f3 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 10:35:35 +0900
Subject: [PATCH 05/77] Add missing tf.keras.Model methods

---
 .../tensorflow/tensorflow/keras/__init__.pyi  |  2 +-
 stubs/tensorflow/tensorflow/keras/models.pyi  | 20 ++++++++++++++++++-
 stubs/tensorflow/tensorflow/train.py          |  8 ++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 stubs/tensorflow/tensorflow/train.py

diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 5ea092c24306..4f7d70cc36f6 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -7,9 +7,9 @@ from tensorflow.keras import (
     layers as layers,
     losses as losses,
     metrics as metrics,
+    models as models,
     optimizers as optimizers,
     regularizers as regularizers,
-    models as models,
 )
 
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 787563f5c769..f38c2891c76d 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,3 +1,21 @@
+from collections.abc import Callable
+from pathlib import Path
+
+_ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
+
 # class Model(base_layer.Layer, version_utils.ModelVersionSelector):
-class Model():
+class Model:
     def __init__(self, *args, **kwargs) -> None: ...
+    def build(self, input_shape: _ShapeLike): ...
+    def summary(
+        self,
+        line_length: None | int = None,
+        positions: None | list[float] = None,
+        print_fn: None | Callable[[str], None] = None,
+        expand_nested: bool = False,
+        show_trainable: bool = False,
+        layer_range: None | list[str] | tuple[str, str] =None,
+    ): ...
+    def load_weights(
+        self, filepath: str | Path, skip_mismatch: bool = False, by_name: bool = False, options: None | tensorflow.train.CheckpointOptions =None
+    ): ...
diff --git a/stubs/tensorflow/tensorflow/train.py b/stubs/tensorflow/tensorflow/train.py
new file mode 100644
index 000000000000..737ffe5e8c6c
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/train.py
@@ -0,0 +1,8 @@
+from collections.abc import Callable
+
+
+class CheckpointOptions:
+    experimental_io_device: None | str = None
+    experimental_enable_async_checkpoint: bool = False
+    experimental_write_callbacks: None | list[Callable] = None
+    enable_async: bool = False

From cbc264171a97f49ae30495f0fec86e7c864e5cf0 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 17:49:35 +0900
Subject: [PATCH 06/77] Add TF2 functions/models/methods.

---
 stubs/tensorflow/tensorflow/__init__.pyi      | 19 +++++--
 .../tensorflow/tensorflow/keras/__init__.pyi  |  1 +
 stubs/tensorflow/tensorflow/keras/layers.pyi  | 55 ++++++++++++++++++-
 stubs/tensorflow/tensorflow/keras/models.pyi  | 13 +++--
 stubs/tensorflow/tensorflow/math.pyi          |  6 ++
 5 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 5b0e66671c25..72385c7127dd 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -73,7 +73,7 @@ from tensorflow.sparse import SparseTensor as SparseTensor
 # we will skip making Tensor generic. Also good type hints for shapes will
 # run quickly into many places where type system is not strong enough today.
 # So shape typing is probably not worth doing anytime soon.
-_Slice: TypeAlias = int | slice | None
+_Slice: TypeAlias = int | slice | Tensor | None
 
 _FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
 _StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
@@ -219,19 +219,20 @@ class TensorShape(metaclass=ABCMeta):
     def __init__(self, dims: _ShapeLike) -> None: ...
     @property
     def rank(self) -> int: ...
-    def as_list(self) -> list[int | None]: ...
+    def as_list(self) -> list[int]: ...
     def assert_has_rank(self, rank: int) -> None: ...
     def assert_is_compatible_with(self, other: Iterable[int | None]) -> None: ...
     def __bool__(self) -> _bool: ...
     @overload
-    def __getitem__(self, key: int) -> int | None: ...
+    def __getitem__(self, key: int) -> int: ...
     @overload
     def __getitem__(self, key: slice) -> TensorShape: ...
-    def __iter__(self) -> Iterator[int | None]: ...
+    def __iter__(self) -> Iterator[int]: ...
     def __len__(self) -> int: ...
     def __add__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __radd__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __getattr__(self, name: str) -> Incomplete: ...
+    def __index__(self) -> int: ...
 
 class Graph:
     def add_to_collection(self, name: str, value: object) -> None: ...
@@ -266,6 +267,10 @@ class IndexedSlices(metaclass=ABCMeta):
     @property
     def device(self) -> str: ...
     def __neg__(self) -> IndexedSlices: ...
+    @overload
+    def __getitem__(self, key: int) -> Tensor: ...
+    @overload
+    def __getitem__(self, key: slice) -> Iterable[Tensor]: ...
     def consumers(self) -> list[Operation]: ...
 
 class name_scope:
@@ -411,5 +416,9 @@ def convert_to_tensor(
 ) -> Tensor: ...
 def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
 # TODO: The "list"s below are from the TF doc, any Iterable or Sequence might be enough ?
-def concat(values: list[Tensor] | Tensor, axis: int, name: None | str = "concat") -> Tensor: ...
+def concat(values: Iterable[Tensor] | Tensor, axis: int, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | list[int] = None, name: None | str = None) -> Tensor: ...
+def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
+def constant(value: Sequence, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
+def cast(x: Tensor | _, dtype, name: DType | str = None) -> Tensor: ...
+def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...
diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 4f7d70cc36f6..3a43597e1c6d 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -11,5 +11,6 @@ from tensorflow.keras import (
     optimizers as optimizers,
     regularizers as regularizers,
 )
+from tensorflow.keras.models import Model as Model
 
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi
index 6c38244574e0..fbc5b1e8623d 100644
--- a/stubs/tensorflow/tensorflow/keras/layers.pyi
+++ b/stubs/tensorflow/tensorflow/keras/layers.pyi
@@ -1,6 +1,6 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Generic, TypeVar, overload
+from typing import Any, Generic, Literal, TypeVar, overload
 from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
@@ -194,4 +194,57 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]):
         name: str | None = None,
     ) -> None: ...
 
+class Conv2D(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self,
+        filters: int,
+        kernel_size: int | tuple[int, int],
+        strides: int | tuple[int, int] = (1, 1),
+        padding: Literal["valid", "same"] = "valid",
+        data_format: None | Literal["channels_last", "channels_first"] = None,
+        dilation_rate: int | tuple[int, int] = (1, 1),
+        groups: int = 1,
+        activation: _Activation = None,
+        use_bias: bool = True,
+        kernel_initializer: _Initializer = "glorot_uniform",
+        bias_initializer: _Initializer = "zeros",
+        kernel_regularizer: _Regularizer = None,
+        bias_regularizer: _Regularizer = None,
+        activity_regularizer: _Regularizer = None,
+        kernel_constraint: _Constraint = None,
+        bias_constraint: _Constraint = None,
+        trainable: bool = True,
+        dtype: _LayerDtype = None,
+        dynamic: bool = False,
+        name: str | None = None,
+    ) -> None: ...
+
+class Identity(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self,
+        trainable: bool = True,
+        dtype: _LayerDtype = None,
+        dynamic: bool = False,
+        name: str | None = None,
+    ) -> None: ...
+
+class LayerNormalization(Layer[tf.Tensor, tf.Tensor]):
+    def __init__(
+        self,
+        axis: int = -1,
+        epsilon: float = 0.001,
+        center: bool = True,
+        scale: bool = True,
+        beta_initializer: _Initializer = "zeros",
+        gamma_initializer: _Initializer = "ones",
+        beta_regularizer: _Regularizer = None,
+        gamma_regularizer: _Regularizer = None,
+        beta_constraint: _Constraint = None,
+        gamma_constraint: _Constraint = None,
+        trainable: bool = True,
+        dtype: _LayerDtype = None,
+        dynamic: bool = False,
+        name: str | None = None,
+    ) -> None: ...
+
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index f38c2891c76d..51bc6d4996af 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,12 +1,15 @@
 from collections.abc import Callable
 from pathlib import Path
+from typing import Generic
 
-_ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
+import tensorflow
+import tensorflow as tf
+from tensorflow import _ShapeLike, _TensorCompatible
+from tensorflow.keras.layers import _InputT, _OutputT
 
-# class Model(base_layer.Layer, version_utils.ModelVersionSelector):
-class Model:
+class Model(tf.Module, tf.keras.layers.Layer, Generic[_InputT, _OutputT]):
     def __init__(self, *args, **kwargs) -> None: ...
-    def build(self, input_shape: _ShapeLike): ...
+    def build(self, input_shape: _ShapeLike) -> None: ...
     def summary(
         self,
         line_length: None | int = None,
@@ -19,3 +22,5 @@ class Model:
     def load_weights(
         self, filepath: str | Path, skip_mismatch: bool = False, by_name: bool = False, options: None | tensorflow.train.CheckpointOptions =None
     ): ...
+    def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def call(self, __inputs: _InputT) -> _OutputT: ...
diff --git a/stubs/tensorflow/tensorflow/math.pyi b/stubs/tensorflow/tensorflow/math.pyi
index 66fc425dbe40..d26f45896efb 100644
--- a/stubs/tensorflow/tensorflow/math.pyi
+++ b/stubs/tensorflow/tensorflow/math.pyi
@@ -252,6 +252,12 @@ def reduce_std(
     keepdims: bool = False,
     name: str | None = None,
 ) -> Tensor: ...
+def reduce_variance(
+    input_tensor: _TensorCompatible | RaggedTensor,
+    axis: _TensorCompatible | None = None,
+    keepdims: bool = False,
+    name: str | None = None,
+) -> Tensor: ...
 def argmax(
     input: _TensorCompatible, axis: _TensorCompatible | None = None, output_type: _DTypeLike = ..., name: str | None = None
 ) -> Tensor: ...

From 5a353e38e61aa47c6e30bc267acd0c9823cae451 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:01:24 +0900
Subject: [PATCH 07/77] revert TensorShape changes.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 72385c7127dd..523aefa94942 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -219,15 +219,15 @@ class TensorShape(metaclass=ABCMeta):
     def __init__(self, dims: _ShapeLike) -> None: ...
     @property
     def rank(self) -> int: ...
-    def as_list(self) -> list[int]: ...
+    def as_list(self) -> list[int | None]: ...
     def assert_has_rank(self, rank: int) -> None: ...
     def assert_is_compatible_with(self, other: Iterable[int | None]) -> None: ...
     def __bool__(self) -> _bool: ...
     @overload
-    def __getitem__(self, key: int) -> int: ...
+    def __getitem__(self, key: int) -> int | None: ...
     @overload
     def __getitem__(self, key: slice) -> TensorShape: ...
-    def __iter__(self) -> Iterator[int]: ...
+    def __iter__(self) -> Iterator[int | None]: ...
     def __len__(self) -> int: ...
     def __add__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __radd__(self, other: Iterable[int | None]) -> TensorShape: ...

From c265a66882615bea480a3ff0e1e410e9adb06302 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:05:18 +0900
Subject: [PATCH 08/77] Fix tf functions.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 523aefa94942..b5f76499e5de 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -415,10 +415,9 @@ def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
 def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
-# TODO: The "list"s below are from the TF doc, any Iterable or Sequence might be enough ?
-def concat(values: Iterable[Tensor] | Tensor, axis: int, name: None | str = "concat") -> Tensor: ...
-def squeeze(input: Tensor, axis: None | list[int] = None, name: None | str = None) -> Tensor: ...
+def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
+def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
-def constant(value: Sequence, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
-def cast(x: Tensor | _, dtype, name: DType | str = None) -> Tensor: ...
+def constant(value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
+def cast(x: Tensor | _TensorCompatible, dtype: DType, name: None | str = None) -> Tensor: ...
 def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...

From eb3098f694c6393e58dc30f5664a761d00d27a43 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:15:40 +0900
Subject: [PATCH 09/77] Fix CheckpointOptions's callbacks type.

---
 stubs/tensorflow/tensorflow/train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/train.py b/stubs/tensorflow/tensorflow/train.py
index 737ffe5e8c6c..dd467d1a0dd4 100644
--- a/stubs/tensorflow/tensorflow/train.py
+++ b/stubs/tensorflow/tensorflow/train.py
@@ -1,8 +1,9 @@
 from collections.abc import Callable
+from typing import Any
 
 
 class CheckpointOptions:
     experimental_io_device: None | str = None
     experimental_enable_async_checkpoint: bool = False
-    experimental_write_callbacks: None | list[Callable] = None
+    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
     enable_async: bool = False

From 005cf22060e6417366e64eb572f5ab0067357541 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 24 Jan 2024 13:20:35 +0000
Subject: [PATCH 10/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/__init__.pyi     | 4 +++-
 stubs/tensorflow/tensorflow/keras/layers.pyi | 6 +-----
 stubs/tensorflow/tensorflow/keras/models.pyi | 8 ++++++--
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index b5f76499e5de..74ed636c67b3 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -418,6 +418,8 @@ def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ..
 def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
-def constant(value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
+def constant(
+    value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = "Const"
+) -> Tensor: ...
 def cast(x: Tensor | _TensorCompatible, dtype: DType, name: None | str = None) -> Tensor: ...
 def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...
diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi
index fbc5b1e8623d..4c29a5d5ec47 100644
--- a/stubs/tensorflow/tensorflow/keras/layers.pyi
+++ b/stubs/tensorflow/tensorflow/keras/layers.pyi
@@ -221,11 +221,7 @@ class Conv2D(Layer[tf.Tensor, tf.Tensor]):
 
 class Identity(Layer[tf.Tensor, tf.Tensor]):
     def __init__(
-        self,
-        trainable: bool = True,
-        dtype: _LayerDtype = None,
-        dynamic: bool = False,
-        name: str | None = None,
+        self, trainable: bool = True, dtype: _LayerDtype = None, dynamic: bool = False, name: str | None = None
     ) -> None: ...
 
 class LayerNormalization(Layer[tf.Tensor, tf.Tensor]):
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 51bc6d4996af..8365d5041a8a 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -17,10 +17,14 @@ class Model(tf.Module, tf.keras.layers.Layer, Generic[_InputT, _OutputT]):
         print_fn: None | Callable[[str], None] = None,
         expand_nested: bool = False,
         show_trainable: bool = False,
-        layer_range: None | list[str] | tuple[str, str] =None,
+        layer_range: None | list[str] | tuple[str, str] = None,
     ): ...
     def load_weights(
-        self, filepath: str | Path, skip_mismatch: bool = False, by_name: bool = False, options: None | tensorflow.train.CheckpointOptions =None
+        self,
+        filepath: str | Path,
+        skip_mismatch: bool = False,
+        by_name: bool = False,
+        options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, __inputs: _InputT) -> _OutputT: ...

From 48a2479afdc95aa99b33e7343eb092976345585f Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:22:04 +0900
Subject: [PATCH 11/77] fix: file name train.py -> train.pyi

---
 stubs/tensorflow/tensorflow/{train.py => train.pyi} | 1 -
 1 file changed, 1 deletion(-)
 rename stubs/tensorflow/tensorflow/{train.py => train.pyi} (99%)

diff --git a/stubs/tensorflow/tensorflow/train.py b/stubs/tensorflow/tensorflow/train.pyi
similarity index 99%
rename from stubs/tensorflow/tensorflow/train.py
rename to stubs/tensorflow/tensorflow/train.pyi
index dd467d1a0dd4..49a78be25b98 100644
--- a/stubs/tensorflow/tensorflow/train.py
+++ b/stubs/tensorflow/tensorflow/train.pyi
@@ -1,7 +1,6 @@
 from collections.abc import Callable
 from typing import Any
 
-
 class CheckpointOptions:
     experimental_io_device: None | str = None
     experimental_enable_async_checkpoint: bool = False

From e1515cc1c6c612a37c1fb236315a11fbcd8137a4 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:30:45 +0900
Subject: [PATCH 12/77] fix: Model MRO error

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 8365d5041a8a..e0c4ccf4c099 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,13 +1,12 @@
 from collections.abc import Callable
 from pathlib import Path
-from typing import Generic
 
 import tensorflow
 import tensorflow as tf
 from tensorflow import _ShapeLike, _TensorCompatible
-from tensorflow.keras.layers import _InputT, _OutputT
+from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
-class Model(tf.Module, tf.keras.layers.Layer, Generic[_InputT, _OutputT]):
+class Model(Layer[_InputT, _OutputT], tf.Module):
     def __init__(self, *args, **kwargs) -> None: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def summary(

From d598706932f5dbf803e94c54dc650532e259af47 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:33:33 +0900
Subject: [PATCH 13/77] fix: concat's values type

---
 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 74ed636c67b3..197dd9c29436 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -415,7 +415,7 @@ def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
 def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
-def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
+def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
 def constant(

From e3911eadcc473eea853c90b0d8ee52f42bb12e20 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:41:25 +0900
Subject: [PATCH 14/77] fix: remove __getitem__ from IndexedSlices.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 197dd9c29436..c665d9bd501a 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -267,10 +267,6 @@ class IndexedSlices(metaclass=ABCMeta):
     @property
     def device(self) -> str: ...
     def __neg__(self) -> IndexedSlices: ...
-    @overload
-    def __getitem__(self, key: int) -> Tensor: ...
-    @overload
-    def __getitem__(self, key: slice) -> Iterable[Tensor]: ...
     def consumers(self) -> list[Operation]: ...
 
 class name_scope:

From f0cc8ebf64b335add2a73b3844a7bf7e8a63e69b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:45:31 +0900
Subject: [PATCH 15/77] fix: remove __index__ from TensorShape.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 1 -
 1 file changed, 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index c665d9bd501a..9b5af4b01137 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -232,7 +232,6 @@ class TensorShape(metaclass=ABCMeta):
     def __add__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __radd__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __getattr__(self, name: str) -> Incomplete: ...
-    def __index__(self) -> int: ...
 
 class Graph:
     def add_to_collection(self, name: str, value: object) -> None: ...

From 2db437a744cdc15dc0077ef2283f2f87f42d9e53 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:53:01 +0900
Subject: [PATCH 16/77] fix: reorganize config module.

---
 .../tensorflow/tensorflow/{config.pyi => config/__init__.pyi}  | 3 ++-
 stubs/tensorflow/tensorflow/config/experimental.pyi            | 3 +++
 2 files changed, 5 insertions(+), 1 deletion(-)
 rename stubs/tensorflow/tensorflow/{config.pyi => config/__init__.pyi} (79%)
 create mode 100644 stubs/tensorflow/tensorflow/config/experimental.pyi

diff --git a/stubs/tensorflow/tensorflow/config.pyi b/stubs/tensorflow/tensorflow/config/__init__.pyi
similarity index 79%
rename from stubs/tensorflow/tensorflow/config.pyi
rename to stubs/tensorflow/tensorflow/config/__init__.pyi
index 14f8099bd2b4..9a9f5f91774c 100644
--- a/stubs/tensorflow/tensorflow/config.pyi
+++ b/stubs/tensorflow/tensorflow/config/__init__.pyi
@@ -1,9 +1,10 @@
 from typing import NamedTuple
 
+from tensorflow.config import experimental as experimental
+
 class PhysicalDevice(NamedTuple):
     name: str
     device_type: str
 
 def list_physical_devices(device_type: None | str = None) -> list[PhysicalDevice]: ...
-def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...
 def set_visible_devices(devices: list[PhysicalDevice], device_type: None | str = None) -> None: ...
diff --git a/stubs/tensorflow/tensorflow/config/experimental.pyi b/stubs/tensorflow/tensorflow/config/experimental.pyi
new file mode 100644
index 000000000000..0cd394802962
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/config/experimental.pyi
@@ -0,0 +1,3 @@
+from tensorflow.config import PhysicalDevice
+
+def set_memory_growth(device: PhysicalDevice, enable: bool) -> None: ...

From 8e011706e8804a53bca202003f1f97a7eb44eb32 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:42:30 +0900
Subject: [PATCH 17/77] fix: add missing Model methods/properties (wip)

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 210 ++++++++++++++++++-
 1 file changed, 200 insertions(+), 10 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index e0c4ccf4c099..41605d45e13c 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,22 +1,187 @@
+from _typeshed import Incomplete
 from collections.abc import Callable
 from pathlib import Path
+from typing import Any, Iterator, Literal, Self
+
+# import numpy as np
+# import numpy.typing as npt
 
 import tensorflow
 import tensorflow as tf
-from tensorflow import _ShapeLike, _TensorCompatible
+from tensorflow import _ShapeLike, _TensorCompatible, Variable
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
+    def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args, **kwargs) -> None: ...
+    def __setattr__(self, name, value) -> None: ...
+    def __reduce__(self) -> Incomplete: ...
+    def __deepcopy__(self, memo) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
-    def summary(
+    def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def compile(
         self,
-        line_length: None | int = None,
-        positions: None | list[float] = None,
-        print_fn: None | Callable[[str], None] = None,
-        expand_nested: bool = False,
-        show_trainable: bool = False,
-        layer_range: None | list[str] | tuple[str, str] = None,
+        optimizer="rmsprop",
+        loss=None,
+        metrics=None,
+        loss_weights=None,
+        weighted_metrics=None,
+        run_eagerly=None,
+        steps_per_execution=None,
+        jit_compile=None,
+        pss_evaluation_shards=0,
+        **kwargs,
+    ) -> Incomplete: ...
+    @property
+    def metrics(self) -> list[Incomplete]: ...
+    @property
+    def metrics_names(self) -> list[str]: ...
+    @property
+    def distribute_strategy(self) -> Incomplete: ...  # tf.distribute.Strategy
+    @property
+    def run_eagerly(self) -> bool: ...
+    @property
+    def autotune_steps_per_execution(self) -> Incomplete: ...
+    @property
+    def steps_per_execution(self) -> int: ...
+    @property
+    def jit_compile(self) -> bool: ...
+    @property
+    def distribute_reduction_method(self) -> Incomplete | Literal["auto"]: ...
+    def train_step(self, data: _TensorCompatible) -> Incomplete: ...
+    def compute_loss(
+        self,
+        x: _TensorCompatible | None = None,
+        y: _TensorCompatible | None = None,
+        y_pred: _TensorCompatible | None = None,
+        sample_weight: Incomplete | None = None,
+    ) -> tf.Tensor | None: ...
+    def compute_metrics(
+        self,
+        x: _TensorCompatible,
+        y: _TensorCompatible,
+        y_pred: _TensorCompatible,
+        sample_weight: Incomplete,
+    ) -> dict[str, float]: ...
+    def get_metrics_result(self) -> dict[str, float]: ...
+    def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
+    def fit(
+        self,
+        x=None,
+        y=None,
+        batch_size=None,
+        epochs=1,
+        verbose="auto",
+        callbacks=None,
+        validation_split=0.0,
+        validation_data=None,
+        shuffle=True,
+        class_weight=None,
+        sample_weight=None,
+        initial_epoch=0,
+        steps_per_epoch=None,
+        validation_steps=None,
+        validation_batch_size=None,
+        validation_freq=1,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+    ): ...
+    def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
+    def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
+    def evaluate(
+        self,
+        x=None,
+        y=None,
+        batch_size=None,
+        verbose="auto",
+        sample_weight=None,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        return_dict=False,
+        **kwargs,
+    ): ...
+    def predict_step(self, data: _InputT) -> _OutputT: ...
+    def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
+    def predict(
+        self,
+        x,
+        batch_size=None,
+        verbose="auto",
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+    ): ...
+    def reset_metrics(self) -> None: ...
+    def train_on_batch(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        class_weight=None,
+        reset_metrics=True,
+        return_dict=False,
+    ) -> float | list[float]: ...
+    def test_on_batch(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        reset_metrics=True,
+        return_dict=False,
+    ) -> float | list[float]: ...
+    def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
+    def fit_generator(
+        self,
+        generator,
+        steps_per_epoch=None,
+        epochs=1,
+        verbose=1,
+        callbacks=None,
+        validation_data=None,
+        validation_steps=None,
+        validation_freq=1,
+        class_weight=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        shuffle=True,
+        initial_epoch=0,
+    ): ...
+    def evaluate_generator(
+        self,
+        generator,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        verbose=0,
+    ): ...
+    def predict_generator(
+        self,
+        generator,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        verbose=0,
+    ): ...
+    @property
+    def trainable_weights(self) -> list[Variable]: ...
+    @property
+    def non_trainable_weights(self) -> list[Variable]: ...
+    def get_weights(self): ...
+    def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
+    def save_weights(
+        self, filepath, overwrite=True, save_format=None, options=None
     ): ...
     def load_weights(
         self,
@@ -25,5 +190,30 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
-    def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    def call(self, __inputs: _InputT) -> _OutputT: ...
+    def get_config(self): ...
+    @classmethod
+    def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
+    def to_json(self, **kwargs) -> str: ...
+    def to_yaml(self, **kwargs) -> str: ...
+    def reset_states(self) -> None: ...
+    @property
+    def state_updates(self) -> list[Incomplete]: ...
+    @property
+    def weights(self) -> list[Variable]: ...
+    def summary(
+        self,
+        line_length: None | int = None,
+        positions: None | list[float] = None,
+        print_fn: None | Callable[[str], None] = None,
+        expand_nested: bool = False,
+        show_trainable: bool = False,
+        layer_range: None | list[str] | tuple[str, str] = None,
+    ): ...
+    @property
+    def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
+    def get_layer(self, name=None, index=None) -> Layer[Incomplete, Incomplete]: ...
+    def get_weight_paths(self): ...
+    def get_compile_config(self) -> dict[str, Any]: ...
+    def compile_from_config(self, config: dict[str, Any]) -> Self: ...
+    def export(self, filepath: str | Path) -> None: ...
+    def save_spec(self, dynamic_batch: bool = True) -> tuple[tuple[tf.TensorSpec, ...], dict[str, tf.TensorSpec]] | None: ...

From 8962c9356dce9f0f3fc9e139f07a4e9cdf04364d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 24 Jan 2024 14:47:15 +0000
Subject: [PATCH 18/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 48 +++-----------------
 1 file changed, 7 insertions(+), 41 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 41605d45e13c..80bcef91637a 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -5,10 +5,9 @@ from typing import Any, Iterator, Literal, Self
 
 # import numpy as np
 # import numpy.typing as npt
-
 import tensorflow
 import tensorflow as tf
-from tensorflow import _ShapeLike, _TensorCompatible, Variable
+from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
@@ -58,11 +57,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         sample_weight: Incomplete | None = None,
     ) -> tf.Tensor | None: ...
     def compute_metrics(
-        self,
-        x: _TensorCompatible,
-        y: _TensorCompatible,
-        y_pred: _TensorCompatible,
-        sample_weight: Incomplete,
+        self, x: _TensorCompatible, y: _TensorCompatible, y_pred: _TensorCompatible, sample_weight: Incomplete
     ) -> dict[str, float]: ...
     def get_metrics_result(self) -> dict[str, float]: ...
     def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
@@ -120,22 +115,9 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ): ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        class_weight=None,
-        reset_metrics=True,
-        return_dict=False,
-    ) -> float | list[float]: ...
-    def test_on_batch(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        reset_metrics=True,
-        return_dict=False,
+        self, x, y=None, sample_weight=None, class_weight=None, reset_metrics=True, return_dict=False
     ) -> float | list[float]: ...
+    def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
@@ -155,24 +137,10 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         initial_epoch=0,
     ): ...
     def evaluate_generator(
-        self,
-        generator,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        verbose=0,
+        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
     ): ...
     def predict_generator(
-        self,
-        generator,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        verbose=0,
+        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
     ): ...
     @property
     def trainable_weights(self) -> list[Variable]: ...
@@ -180,9 +148,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
-    def save_weights(
-        self, filepath, overwrite=True, save_format=None, options=None
-    ): ...
+    def save_weights(self, filepath, overwrite=True, save_format=None, options=None): ...
     def load_weights(
         self,
         filepath: str | Path,

From 8f7df4688063d2b2d0f1e7ba9655adc1199fb8dc Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:50:50 +0900
Subject: [PATCH 19/77] fix: use Iterator from collections.abc instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 80bcef91637a..0d642f273cfc 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Any, Iterator, Literal, Self
+from typing import Any, Literal, Self
 
 # import numpy as np
 # import numpy.typing as npt

From 8bc143fe2e066ccdd17de04ba6531b2d7a114327 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:51:31 +0900
Subject: [PATCH 20/77] fix: import Self from typing_extensions instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 0d642f273cfc..4b27b8272450 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,7 +1,8 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Any, Literal, Self
+from typing import Any, Literal
+from typing_extensions import Self
 
 # import numpy as np
 # import numpy.typing as npt

From 3e276248b57ef8263c42f65d4cb41b6c275b2643 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:13:49 +0900
Subject: [PATCH 21/77] feat: add empty Metric class

---
 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 de1c1f79f665..9f57d524f34d 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,5 +1,9 @@
+from _typeshed import Incomplete
 from tensorflow import Tensor, _TensorCompatible
 
+class Metric(Incomplete):
+    ...
+
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1
 ) -> Tensor: ...

From 59f6f6f68b3693c122f12d8cec9ecdfcf0365b4d Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:17:11 +0900
Subject: [PATCH 22/77] feat: Add type hints for keras Model's compile, fit and
 evaluate

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 93 ++++++++++----------
 1 file changed, 47 insertions(+), 46 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 4b27b8272450..ace5efeec42e 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,15 +1,16 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Container
 from pathlib import Path
 from typing import Any, Literal
 from typing_extensions import Self
 
-# import numpy as np
-# import numpy.typing as npt
+import numpy as np
+import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
+from tensorflow.keras.optimizers.legacy import Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
     def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
@@ -22,17 +23,17 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def compile(
         self,
-        optimizer="rmsprop",
-        loss=None,
-        metrics=None,
-        loss_weights=None,
-        weighted_metrics=None,
-        run_eagerly=None,
-        steps_per_execution=None,
-        jit_compile=None,
-        pss_evaluation_shards=0,
+        optimizer: Optimizer | str = "rmsprop",
+        loss: tf.keras.losses.Loss | str | None = None,
+        metrics: list[tf.keras.metrics.Metric | str] | None = None,
+        loss_weights: list[float] | dict[str, float] | None = None,
+        weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
+        run_eagerly: bool | None = None,
+        steps_per_execution: int |  Literal["auto"] | None= None,
+        jit_compile: bool | None = None,
+        pss_evaluation_shards: int |  Literal["auto"] = 0,
         **kwargs,
-    ) -> Incomplete: ...
+    ) -> None: ...
     @property
     def metrics(self) -> list[Incomplete]: ...
     @property
@@ -64,52 +65,52 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def fit(
         self,
-        x=None,
-        y=None,
-        batch_size=None,
-        epochs=1,
-        verbose="auto",
-        callbacks=None,
-        validation_split=0.0,
-        validation_data=None,
-        shuffle=True,
-        class_weight=None,
-        sample_weight=None,
-        initial_epoch=0,
-        steps_per_epoch=None,
-        validation_steps=None,
-        validation_batch_size=None,
-        validation_freq=1,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
+        epochs: int = 1,
+        verbose: Literal["auto", 0, 1, 2] = "auto",
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        validation_split: float = 0.0,
+        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
+        shuffle: bool = True,
+        class_weight: dict[int, float] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        initial_epoch: int = 0,
+        steps_per_epoch: int | None = None,
+        validation_steps: int | None = None,
+        validation_batch_size: int | None = None,
+        validation_freq: int | Container[int] = 1,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
     ): ...
     def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
     def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def evaluate(
         self,
-        x=None,
-        y=None,
-        batch_size=None,
-        verbose="auto",
-        sample_weight=None,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        return_dict=False,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
+        verbose: Literal["auto", 0, 1, 2] = "auto",
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        return_dict: bool = False,
         **kwargs,
     ): ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
         batch_size=None,
-        verbose="auto",
+        verbose: Literal["auto", 0, 1, 2] = "auto",
         steps=None,
-        callbacks=None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
         max_queue_size=10,
         workers=1,
         use_multiprocessing=False,

From 1f216a440b7e134dcabece4e1acc8d0d9b6b8d75 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:31:08 +0900
Subject: [PATCH 23/77] feat: Add type hints for keras Model's predict,
 train_on_batch, fit_generator, evaluate_generator and predict_generator

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 66 +++++++++++++-------
 1 file changed, 43 insertions(+), 23 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index ace5efeec42e..1e6a3a0f21ca 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -106,43 +106,63 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        batch_size=None,
+        x: _TensorCompatible | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
         verbose: Literal["auto", 0, 1, 2] = "auto",
-        steps=None,
+        steps: int | None = None,
         callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
     ): ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
-        self, x, y=None, sample_weight=None, class_weight=None, reset_metrics=True, return_dict=False
+        self,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        class_weight: dict[int, float] | None = None,
+        reset_metrics: bool = True,
+        return_dict: bool = False
     ) -> float | list[float]: ...
     def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
-        generator,
-        steps_per_epoch=None,
-        epochs=1,
-        verbose=1,
-        callbacks=None,
-        validation_data=None,
-        validation_steps=None,
-        validation_freq=1,
-        class_weight=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        shuffle=True,
-        initial_epoch=0,
+        generator: Iterator[Incomplete],
+        steps_per_epoch: int | None = None,
+        epochs: int = 1,
+        verbose: Literal["auto", 0, 1, 2] = 1,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
+        validation_steps: int | None = None,
+        validation_freq: int | Container[int] = 1,
+        class_weight: dict[int, float] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        shuffle: bool = True,
+        initial_epoch: int = 0,
     ): ...
     def evaluate_generator(
-        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
+        self,
+        generator: Iterator[Incomplete],
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        verbose: Literal["auto", 0, 1, 2] = 0,
     ): ...
     def predict_generator(
-        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
+        self,
+        generator: Iterator[Incomplete],
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        verbose: Literal["auto", 0, 1, 2] = 0,
     ): ...
     @property
     def trainable_weights(self) -> list[Variable]: ...

From a6f7ba291039c621a48c5924e86e32c1062faaa8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:31:59 +0000
Subject: [PATCH 24/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 4 ++--
 stubs/tensorflow/tensorflow/keras/models.pyi  | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index 9f57d524f34d..93cb0836fc8d 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,8 +1,8 @@
 from _typeshed import Incomplete
+
 from tensorflow import Tensor, _TensorCompatible
 
-class Metric(Incomplete):
-    ...
+class Metric(Incomplete): ...
 
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 1e6a3a0f21ca..9dd8b6bb5be3 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,5 +1,5 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Iterator, Container
+from collections.abc import Callable, Container, Iterator
 from pathlib import Path
 from typing import Any, Literal
 from typing_extensions import Self
@@ -29,9 +29,9 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         loss_weights: list[float] | dict[str, float] | None = None,
         weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
         run_eagerly: bool | None = None,
-        steps_per_execution: int |  Literal["auto"] | None= None,
+        steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
-        pss_evaluation_shards: int |  Literal["auto"] = 0,
+        pss_evaluation_shards: int | Literal["auto"] = 0,
         **kwargs,
     ) -> None: ...
     @property
@@ -123,7 +123,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         sample_weight: npt.NDArray[np.float_] | None = None,
         class_weight: dict[int, float] | None = None,
         reset_metrics: bool = True,
-        return_dict: bool = False
+        return_dict: bool = False,
     ) -> float | list[float]: ...
     def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]

From f7cc67dcc0432c50c2a109456d0e09fc19c73a0c Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:35:56 +0900
Subject: [PATCH 25/77] feat: Add type hints for keras Model's save,
 save_weights and load_weights

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 9dd8b6bb5be3..84c268ae5aad 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -169,8 +169,20 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
-    def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
-    def save_weights(self, filepath, overwrite=True, save_format=None, options=None): ...
+    def save(
+        self,
+        filepath: str | Path,
+        overwrite: bool = True,
+        save_format: Literal["keras", "tf", "h5"] | None = None,
+        **kwargs,
+    ): ...
+    def save_weights(
+        self,
+        filepath: str | Path,
+        overwrite: bool = True,
+        save_format: Literal["tf", "h5"] | None = None,
+        options: tf.train.CheckpointOptions | None = None,
+    ): ...
     def load_weights(
         self,
         filepath: str | Path,
@@ -178,7 +190,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
-    def get_config(self): ...
+    def get_config(self) -> None: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
     def to_json(self, **kwargs) -> str: ...

From e3573cf7ec5fb85529c48d3a828f10acfc6f53bb Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:37:11 +0900
Subject: [PATCH 26/77] feat: Add type hints for keras Model's get_layer

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 84c268ae5aad..7dbe3a9b858a 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -211,7 +211,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ): ...
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
-    def get_layer(self, name=None, index=None) -> Layer[Incomplete, Incomplete]: ...
+    def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
     def get_weight_paths(self): ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...

From f31992b6824c9b7c8c8538cd5b8d52c4f053f05c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:38:10 +0000
Subject: [PATCH 27/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7dbe3a9b858a..5b0407548af2 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -170,11 +170,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(
-        self,
-        filepath: str | Path,
-        overwrite: bool = True,
-        save_format: Literal["keras", "tf", "h5"] | None = None,
-        **kwargs,
+        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs
     ): ...
     def save_weights(
         self,

From 6c63b89c2e4a8fd47d03941c1163259341cdfdd8 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:45:36 +0900
Subject: [PATCH 28/77] fix: add Any type to args and kwargs.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 5b0407548af2..bc14e89b9587 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -13,8 +13,8 @@ from tensorflow.keras.layers import Layer, _InputT, _OutputT
 from tensorflow.keras.optimizers.legacy import Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
-    def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
-    def __init__(self, *args, **kwargs) -> None: ...
+    def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def __setattr__(self, name, value) -> None: ...
     def __reduce__(self) -> Incomplete: ...
     def __deepcopy__(self, memo) -> Incomplete: ...
@@ -32,7 +32,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
         pss_evaluation_shards: int | Literal["auto"] = 0,
-        **kwargs,
+        **kwargs: Any,
     ) -> None: ...
     @property
     def metrics(self) -> list[Incomplete]: ...
@@ -100,7 +100,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         return_dict: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ): ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
@@ -170,7 +170,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(
-        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs
+        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
     ): ...
     def save_weights(
         self,
@@ -189,8 +189,8 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def get_config(self) -> None: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
-    def to_json(self, **kwargs) -> str: ...
-    def to_yaml(self, **kwargs) -> str: ...
+    def to_json(self, **kwargs: Any) -> str: ...
+    def to_yaml(self, **kwargs: Any) -> str: ...
     def reset_states(self) -> None: ...
     @property
     def state_updates(self) -> list[Incomplete]: ...

From d7b08a613bc9baa6c2c2716d17962461bc4be061 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:54:01 +0900
Subject: [PATCH 29/77] fix: allow tuple[int, ...] as input to expand_dims

This is not written in t he doc, but is used by jax
---
 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 9b5af4b01137..e3ce0cefdc98 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -409,7 +409,7 @@ def __getattr__(name: str) -> Incomplete: ...
 def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
-def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
+def expand_dims(input: Tensor | tuple[int, ...], axis: int, name: None | str = None) -> Tensor: ...
 def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...

From f10de3f7857d9a0e10d3ebe5f612aea0eeda9fb7 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:58:40 +0900
Subject: [PATCH 30/77] fix: add empty tensorflow.kerras.callbacks.Callback
 calls

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/keras/callbacks.pyi

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
new file mode 100644
index 000000000000..191a5db61f49
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -0,0 +1,4 @@
+from _typeshed import Incomplete
+
+
+class Callback(Incomplete): ...

From 41b5a237df225c479308067d332814115b949747 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:59:26 +0000
Subject: [PATCH 31/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 1 -
 1 file changed, 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 191a5db61f49..7e078b7b4573 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -1,4 +1,3 @@
 from _typeshed import Incomplete
 
-
 class Callback(Incomplete): ...

From 9d1d1b4ebac52ef470708704c66598c2b2a296e1 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:59:35 +0900
Subject: [PATCH 32/77] fix: Model's get_config return type.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index bc14e89b9587..ba3f5839f2b7 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -186,7 +186,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
-    def get_config(self) -> None: ...
+    def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
     def to_json(self, **kwargs: Any) -> str: ...

From b63f23d262ff53c5ef732fd7ce445e2b79c1e7b6 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:13:27 +0900
Subject: [PATCH 33/77] feat: add empty History class

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 1 +
 1 file changed, 1 insertion(+)

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 7e078b7b4573..e7d68d4c32b7 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -1,3 +1,4 @@
 from _typeshed import Incomplete
 
 class Callback(Incomplete): ...
+class History(Incomplete): ...

From 2a422dc9703e8656e70b9f9076cf84638ff4abeb Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:14:02 +0900
Subject: [PATCH 34/77] fix: add type hints for test_on_batch, add missing
 return types

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 36 ++++++++++++--------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index ba3f5839f2b7..7bfe829ae27f 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -15,9 +15,9 @@ from tensorflow.keras.optimizers.legacy import Optimizer
 class Model(Layer[_InputT, _OutputT], tf.Module):
     def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
-    def __setattr__(self, name, value) -> None: ...
+    def __setattr__(self, name: str, value: Any) -> None: ...
     def __reduce__(self) -> Incomplete: ...
-    def __deepcopy__(self, memo) -> Incomplete: ...
+    def __deepcopy__(self, memo: Incomplete) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
@@ -84,7 +84,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         max_queue_size: int = 10,
         workers: int = 1,
         use_multiprocessing: bool = False,
-    ): ...
+    ) -> tf.keras.callbacks.History: ...
     def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
     def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def evaluate(
@@ -101,7 +101,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         use_multiprocessing: bool = False,
         return_dict: bool = False,
         **kwargs: Any,
-    ): ...
+    ) -> float | list[float]: ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
@@ -114,7 +114,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         max_queue_size: int = 10,
         workers: int = 1,
         use_multiprocessing: bool = False,
-    ): ...
+    ) -> _OutputT: ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
         self,
@@ -125,7 +125,15 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         reset_metrics: bool = True,
         return_dict: bool = False,
     ) -> float | list[float]: ...
-    def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
+    def test_on_batch(
+        self,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        class_weight: dict[int, float] | None = None,
+        reset_metrics: bool = True,
+        return_dict: bool = False,
+    ) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
@@ -143,7 +151,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         use_multiprocessing: bool = False,
         shuffle: bool = True,
         initial_epoch: int = 0,
-    ): ...
+    ) -> tf.keras.callbacks.History: ...
     def evaluate_generator(
         self,
         generator: Iterator[Incomplete],
@@ -153,7 +161,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         verbose: Literal["auto", 0, 1, 2] = 0,
-    ): ...
+    ) -> float | list[float]: ...
     def predict_generator(
         self,
         generator: Iterator[Incomplete],
@@ -163,7 +171,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         verbose: Literal["auto", 0, 1, 2] = 0,
-    ): ...
+    ) -> _OutputT: ...
     @property
     def trainable_weights(self) -> list[Variable]: ...
     @property
@@ -171,21 +179,21 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def get_weights(self): ...
     def save(
         self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
-    ): ...
+    ) -> None: ...
     def save_weights(
         self,
         filepath: str | Path,
         overwrite: bool = True,
         save_format: Literal["tf", "h5"] | None = None,
         options: tf.train.CheckpointOptions | None = None,
-    ): ...
+    ) -> None: ...
     def load_weights(
         self,
         filepath: str | Path,
         skip_mismatch: bool = False,
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
-    ): ...
+    ) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
@@ -204,11 +212,11 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         expand_nested: bool = False,
         show_trainable: bool = False,
         layer_range: None | list[str] | tuple[str, str] = None,
-    ): ...
+    ) -> None: ...
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
     def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
-    def get_weight_paths(self): ...
+    def get_weight_paths(self) -> dict[str,  tf.Variable ]: ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...
     def export(self, filepath: str | Path) -> None: ...

From d2f218e6db7cdb9b13a7dd2c25947df0087c6905 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 13:14:45 +0000
Subject: [PATCH 35/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7bfe829ae27f..3273ae2b2470 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -216,7 +216,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
     def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
-    def get_weight_paths(self) -> dict[str,  tf.Variable ]: ...
+    def get_weight_paths(self) -> dict[str, tf.Variable]: ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...
     def export(self, filepath: str | Path) -> None: ...

From 5d9933759537bafd496b927f5ee42efc6cf8c9b2 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:20:27 +0900
Subject: [PATCH 36/77] fix: add missing types

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 3273ae2b2470..f0d0025aa95f 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -176,7 +176,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def trainable_weights(self) -> list[Variable]: ...
     @property
     def non_trainable_weights(self) -> list[Variable]: ...
-    def get_weights(self): ...
+    def get_weights(self) -> Incomplete: ...
     def save(
         self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
     ) -> None: ...
@@ -196,7 +196,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
-    def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
+    def from_config(cls, config: dict[str, Any], custom_objects: Incomplete | None = None) -> Self: ...
     def to_json(self, **kwargs: Any) -> str: ...
     def to_yaml(self, **kwargs: Any) -> str: ...
     def reset_states(self) -> None: ...

From 924a59ffd905c57136eed2ebc9a22dc94478952a Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:29:37 +0900
Subject: [PATCH 37/77] fix: allow call's training to be None

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index f0d0025aa95f..dc825d77fa81 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -20,7 +20,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def __deepcopy__(self, memo: Incomplete) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def compile(
         self,
         optimizer: Optimizer | str = "rmsprop",

From f5384f36264129d3705771b1552df6827337ffb9 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:35:29 +0900
Subject: [PATCH 38/77] fix: predict input type.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index dc825d77fa81..e7a9cc4fb611 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -45,7 +45,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def autotune_steps_per_execution(self) -> Incomplete: ...
     @property
-    def steps_per_execution(self) -> int: ...
+    def steps_per_execution(self) -> int: ...  # Requires a compiled model.
     @property
     def jit_compile(self) -> bool: ...
     @property
@@ -106,7 +106,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x: _TensorCompatible | tf.data.Dataset[Incomplete] | None = None,
+        x: _TensorCompatible | tf.data.Dataset[Incomplete],
         batch_size: int | None = None,
         verbose: Literal["auto", 0, 1, 2] = "auto",
         steps: int | None = None,
@@ -130,7 +130,6 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
         y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
         sample_weight: npt.NDArray[np.float_] | None = None,
-        class_weight: dict[int, float] | None = None,
         reset_metrics: bool = True,
         return_dict: bool = False,
     ) -> float | list[float]: ...

From 29b0d4a2ef01b365497c7fd22315425c30a648be Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:46:56 +0900
Subject: [PATCH 39/77] fix: fix CheckpointOptions (add init)

---
 stubs/tensorflow/tensorflow/train.pyi | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/train.pyi b/stubs/tensorflow/tensorflow/train.pyi
index 49a78be25b98..8afafdf972f0 100644
--- a/stubs/tensorflow/tensorflow/train.pyi
+++ b/stubs/tensorflow/tensorflow/train.pyi
@@ -2,7 +2,15 @@ from collections.abc import Callable
 from typing import Any
 
 class CheckpointOptions:
-    experimental_io_device: None | str = None
-    experimental_enable_async_checkpoint: bool = False
-    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
-    enable_async: bool = False
+    experimental_io_device: None | str
+    experimental_enable_async_checkpoint: bool
+    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]]
+    enable_async: bool
+
+    def __init__(
+        self,
+        experimental_io_device: None | str = None,
+        experimental_enable_async_checkpoint: bool = False,
+        experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None,
+        enable_async: bool = False,
+    ) -> None: ...

From 0326484456bcba8d39f8be381edfe7f33d079290 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:47:39 +0900
Subject: [PATCH 40/77] fix: add __init__ and __new__ to Metric

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index 93cb0836fc8d..e5e474a0e53e 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,8 +1,11 @@
-from _typeshed import Incomplete
+from typing import Any, Self
 
 from tensorflow import Tensor, _TensorCompatible
+from tensorflow.dtypes import DType
 
-class Metric(Incomplete): ...
+class Metric:
+    def __init__(self, name: str, dtype: DType) -> None: ...
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
 
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1

From f0f1c0160441ea066755a5f40be99d2d2bdd9a8e Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:48:07 +0900
Subject: [PATCH 41/77] fix: comment out properties not present at runtime.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index e7a9cc4fb611..a9e3d4799440 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -31,6 +31,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         run_eagerly: bool | None = None,
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
+        *,
         pss_evaluation_shards: int | Literal["auto"] = 0,
         **kwargs: Any,
     ) -> None: ...
@@ -42,10 +43,10 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def distribute_strategy(self) -> Incomplete: ...  # tf.distribute.Strategy
     @property
     def run_eagerly(self) -> bool: ...
-    @property
-    def autotune_steps_per_execution(self) -> Incomplete: ...
-    @property
-    def steps_per_execution(self) -> int: ...  # Requires a compiled model.
+    # @property
+    # def autotune_steps_per_execution(self) -> Incomplete: ...  # not present at runtime
+    # @property
+    # def steps_per_execution(self) -> int: ...  # Requires a compiled model. # not present at runtime
     @property
     def jit_compile(self) -> bool: ...
     @property

From 3c772c5297a3df87c20f8ca2fcc04731536499b9 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 23:14:56 +0900
Subject: [PATCH 42/77] fix: import Self from typing_extensions instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index e5e474a0e53e..f27206828a99 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,4 +1,5 @@
-from typing import Any, Self
+from typing import Any
+from typing_extensions import Self
 
 from tensorflow import Tensor, _TensorCompatible
 from tensorflow.dtypes import DType

From 1e804ab240bd19ca915c3ad96f29fff6abf85e78 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Fri, 26 Jan 2024 16:32:11 +0900
Subject: [PATCH 43/77] adding alias types, wip

---
 stubs/tensorflow/tensorflow/__init__.pyi      | 24 ++++-----
 stubs/tensorflow/tensorflow/_aliases.pyi      | 49 ++++++++++++++-----
 stubs/tensorflow/tensorflow/data/__init__.pyi | 22 ++++-----
 stubs/tensorflow/tensorflow/io/__init__.pyi   |  4 +-
 .../tensorflow/tensorflow/keras/__init__.pyi  | 16 ++++++
 stubs/tensorflow/tensorflow/keras/layers.pyi  |  6 +--
 stubs/tensorflow/tensorflow/keras/losses.pyi  |  4 +-
 stubs/tensorflow/tensorflow/keras/models.pyi  | 22 ++++++---
 .../keras/optimizers/legacy/__init__.pyi      | 10 ++--
 9 files changed, 105 insertions(+), 52 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index e3ce0cefdc98..3fb6c2bde55c 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -19,7 +19,7 @@ from tensorflow import (
     keras as keras,
     math as math,
 )
-from tensorflow._aliases import ContainerGradients, ContainerTensors, ContainerTensorsLike, Gradients, TensorLike
+from tensorflow._aliases import _ContainerGradients, _ContainerTensors, _ContainerTensorsLike, _Gradients, _TensorLike
 from tensorflow.core.protobuf import struct_pb2
 
 # Explicit import of DType is covered by the wildcard, but
@@ -307,39 +307,39 @@ class GradientTape:
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
-        sources: TensorLike,
+        target: _ContainerTensors,
+        sources: _TensorLike,
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> Gradients: ...
+    ) -> _Gradients: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
+        target: _ContainerTensors,
         sources: Sequence[Tensor],
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> list[Gradients]: ...
+    ) -> list[_Gradients]: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
+        target: _ContainerTensors,
         sources: Mapping[str, Tensor],
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> dict[str, Gradients]: ...
+    ) -> dict[str, _Gradients]: ...
     @overload
     def gradient(
         self,
-        target: ContainerTensors,
-        sources: ContainerTensors,
+        target: _ContainerTensors,
+        sources: _ContainerTensors,
         output_gradients: list[Tensor] | None = None,
         unconnected_gradients: UnconnectedGradients = ...,
-    ) -> ContainerGradients: ...
+    ) -> _ContainerGradients: ...
     @contextmanager
     def stop_recording(self) -> Generator[None, None, None]: ...
     def reset(self) -> None: ...
-    def watch(self, tensor: ContainerTensorsLike) -> None: ...
+    def watch(self, tensor: _ContainerTensorsLike) -> None: ...
     def watched_variables(self) -> tuple[Variable, ...]: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi
index 555f31a9df24..bf7d4d3ea1db 100644
--- a/stubs/tensorflow/tensorflow/_aliases.pyi
+++ b/stubs/tensorflow/tensorflow/_aliases.pyi
@@ -5,21 +5,21 @@
 from collections.abc import Mapping, Sequence
 from typing import Any, Protocol, TypeVar
 from typing_extensions import TypeAlias
+from typing import Any, Iterable, Mapping, Sequence, TypeVar
 
-import numpy
+import numpy as np
 import tensorflow as tf
+from tensorflow.keras.layers import InputSpec
 
 _T1 = TypeVar("_T1")
-ContainerGeneric: TypeAlias = Mapping[str, ContainerGeneric[_T1]] | Sequence[ContainerGeneric[_T1]] | _T1
+_T = TypeVar("_T")
+_ContainerGeneric: TypeAlias = Mapping[str, _ContainerGeneric[_T1]] | Sequence[_ContainerGeneric[_T1]] | _T1
 
-TensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor | tf.SparseTensor
-Gradients: TypeAlias = tf.Tensor | tf.IndexedSlices
-
-ContainerTensorsLike: TypeAlias = ContainerGeneric[TensorLike]
-ContainerTensors: TypeAlias = ContainerGeneric[tf.Tensor]
-ContainerGradients: TypeAlias = ContainerGeneric[Gradients]
-
-AnyArray: TypeAlias = numpy.ndarray[Any, Any]
+_TensorLike: TypeAlias = tf.Tensor | tf.RaggedTensor | tf.SparseTensor
+_SparseTensorLike = tf.Tensor | tf.SparseTensor
+_RaggedTensorLike = tf.Tensor | tf.RaggedTensor
+_RaggedTensorLikeT = TypeVar("_RaggedTensorLikeT", tf.Tensor, tf.RaggedTensor)
+_Gradients: TypeAlias = tf.Tensor | tf.IndexedSlices
 
 class _KerasSerializable1(Protocol):
     def get_config(self) -> dict[str, Any]: ...
@@ -27,4 +27,31 @@ class _KerasSerializable1(Protocol):
 class _KerasSerializable2(Protocol):
     __name__: str
 
-KerasSerializable: TypeAlias = _KerasSerializable1 | _KerasSerializable2
+_KerasSerializable: TypeAlias = _KerasSerializable1 | _KerasSerializable2
+
+_FloatDataSequence = Sequence[float] | Sequence[_FloatDataSequence]
+_StrDataSequence = Sequence[str] | Sequence[_StrDataSequence]
+_ScalarTensorCompatible = tf.Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+
+_TensorCompatible = _ScalarTensorCompatible | Sequence[_TensorCompatible]
+_TensorCompatibleT = TypeVar("_TensorCompatibleT", bound=_TensorCompatible)
+# Sparse tensors are very annoying. Some operations work on them, but many do not. You
+# will need to manually verify if an operation supports them. SparseTensorCompatible is intended to be a
+# broader type than TensorCompatible and not all operations will support broader version. If unsure,
+# use TensorCompatible instead.
+_SparseTensorCompatible = _TensorCompatible | tf.SparseTensor
+
+_ShapeLike = tf.TensorShape | Iterable[_ScalarTensorCompatible | None] | int | tf.Tensor
+_DTypeLike = tf.DType | str | np.dtype[Any] | int
+_GradientsT = tf.Tensor | tf.IndexedSlices
+
+_ContainerTensors = _ContainerGeneric[tf.Tensor]
+_ContainerTensorsLike = _ContainerGeneric[_TensorLike]
+_ContainerTensorCompatible = _ContainerGeneric[_TensorCompatible]
+_ContainerGradients = _ContainerGeneric[_GradientsT]
+_ContainerTensorShape = _ContainerGeneric[tf.TensorShape]
+_ContainerInputSpec = _ContainerGeneric[InputSpec]
+
+_AnyArray = np.ndarray[Any, Any]
+_FloatArray = np.ndarray[Any, np.dtype[np.float_ | np.float16 | np.float32 | np.float64]]
+_IntArray = np.ndarray[Any, np.dtype[np.int_ | np.uint8 | np.int32 | np.int64]]
diff --git a/stubs/tensorflow/tensorflow/data/__init__.pyi b/stubs/tensorflow/tensorflow/data/__init__.pyi
index 7567aaf07bab..aacb6bb258ed 100644
--- a/stubs/tensorflow/tensorflow/data/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/data/__init__.pyi
@@ -7,7 +7,7 @@ from typing_extensions import Self
 import numpy as np
 import tensorflow as tf
 from tensorflow import TypeSpec, _ScalarTensorCompatible, _TensorCompatible
-from tensorflow._aliases import ContainerGeneric
+from tensorflow._aliases import _ContainerGeneric
 from tensorflow.data import experimental as experimental
 from tensorflow.data.experimental import AUTOTUNE as AUTOTUNE
 from tensorflow.dtypes import DType
@@ -21,7 +21,7 @@ _T3 = TypeVar("_T3")
 class Iterator(_Iterator[_T1], Trackable, ABC):
     @property
     @abstractmethod
-    def element_spec(self) -> ContainerGeneric[TypeSpec[Any]]: ...
+    def element_spec(self) -> _ContainerGeneric[TypeSpec[Any]]: ...
     @abstractmethod
     def get_next(self) -> _T1: ...
     @abstractmethod
@@ -43,8 +43,8 @@ class Dataset(ABC, Generic[_T1]):
         element_length_func: Callable[[_T1], _ScalarTensorCompatible],
         bucket_boundaries: Sequence[int],
         bucket_batch_sizes: Sequence[int],
-        padded_shapes: ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
-        padding_values: ContainerGeneric[_ScalarTensorCompatible] | None = None,
+        padded_shapes: _ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
+        padding_values: _ContainerGeneric[_ScalarTensorCompatible] | None = None,
         pad_to_bucket_boundary: bool = False,
         no_padding: bool = False,
         drop_remainder: bool = False,
@@ -63,7 +63,7 @@ class Dataset(ABC, Generic[_T1]):
     ) -> Dataset[tf.Tensor]: ...
     @property
     @abstractmethod
-    def element_spec(self) -> ContainerGeneric[TypeSpec[Any]]: ...
+    def element_spec(self) -> _ContainerGeneric[TypeSpec[Any]]: ...
     def enumerate(self, start: _ScalarTensorCompatible = 0, name: str | None = None) -> Dataset[tuple[int, _T1]]: ...
     def filter(self, predicate: Callable[[_T1], bool | tf.Tensor], name: str | None = None) -> Dataset[_T1]: ...
     def flat_map(self, map_func: Callable[[_T1], Dataset[_T2]], name: str | None = None) -> Dataset[_T2]: ...
@@ -71,10 +71,10 @@ class Dataset(ABC, Generic[_T1]):
     @staticmethod
     def from_generator(
         generator: Callable[..., _T2],
-        output_types: ContainerGeneric[DType] | None = None,
-        output_shapes: ContainerGeneric[tf.TensorShape | Sequence[int | None]] | None = None,
+        output_types: _ContainerGeneric[DType] | None = None,
+        output_shapes: _ContainerGeneric[tf.TensorShape | Sequence[int | None]] | None = None,
         args: tuple[object, ...] | None = None,
-        output_signature: ContainerGeneric[TypeSpec[Any]] | None = None,
+        output_signature: _ContainerGeneric[TypeSpec[Any]] | None = None,
         name: str | None = None,
     ) -> Dataset[_T2]: ...
     @staticmethod
@@ -111,7 +111,7 @@ class Dataset(ABC, Generic[_T1]):
     @staticmethod
     def load(
         path: str,
-        element_spec: ContainerGeneric[tf.TypeSpec[Any]] | None = None,
+        element_spec: _ContainerGeneric[tf.TypeSpec[Any]] | None = None,
         compression: _CompressionTypes = None,
         reader_func: Callable[[Dataset[Dataset[Any]]], Dataset[Any]] | None = None,
     ) -> Dataset[Any]: ...
@@ -127,8 +127,8 @@ class Dataset(ABC, Generic[_T1]):
     def padded_batch(
         self,
         batch_size: _ScalarTensorCompatible,
-        padded_shapes: ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
-        padding_values: ContainerGeneric[_ScalarTensorCompatible] | None = None,
+        padded_shapes: _ContainerGeneric[tf.TensorShape | _TensorCompatible] | None = None,
+        padding_values: _ContainerGeneric[_ScalarTensorCompatible] | None = None,
         drop_remainder: bool = False,
         name: str | None = None,
     ) -> Dataset[_T1]: ...
diff --git a/stubs/tensorflow/tensorflow/io/__init__.pyi b/stubs/tensorflow/tensorflow/io/__init__.pyi
index 172c79752ecd..bd1ed1389e26 100644
--- a/stubs/tensorflow/tensorflow/io/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/io/__init__.pyi
@@ -5,7 +5,7 @@ from typing import Literal, NamedTuple
 from typing_extensions import Self, TypeAlias
 
 from tensorflow import _DTypeLike, _ShapeLike, _TensorCompatible
-from tensorflow._aliases import TensorLike
+from tensorflow._aliases import _TensorLike
 from tensorflow.io import gfile as gfile
 
 _FeatureSpecs: TypeAlias = Mapping[str, FixedLenFeature | FixedLenSequenceFeature | VarLenFeature | RaggedFeature | SparseFeature]
@@ -102,5 +102,5 @@ class RaggedFeature(NamedTuple):
 
 def parse_example(
     serialized: _TensorCompatible, features: _FeatureSpecs, example_names: Iterable[str] | None = None, name: str | None = None
-) -> dict[str, TensorLike]: ...
+) -> dict[str, _TensorLike]: ...
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 3a43597e1c6d..4705b3f2b1ea 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from typing import Callable
 
 from tensorflow.keras import (
     activations as activations,
@@ -12,5 +13,20 @@ from tensorflow.keras import (
     regularizers as regularizers,
 )
 from tensorflow.keras.models import Model as Model
+import tensorflow as tf
+from tensorflow._aliases import _TensorCompatible
 
 def __getattr__(name: str) -> Incomplete: ...
+
+_Loss = (
+   str
+    | tf.keras.losses.Loss
+    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
+)
+
+_Metric = (
+    str
+    | tf.keras.metrics.Metric
+    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
+    | None
+)
diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi
index 4c29a5d5ec47..d07c0eda7517 100644
--- a/stubs/tensorflow/tensorflow/keras/layers.pyi
+++ b/stubs/tensorflow/tensorflow/keras/layers.pyi
@@ -5,7 +5,7 @@ from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
 from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible
-from tensorflow._aliases import AnyArray
+from tensorflow._aliases import _AnyArray
 from tensorflow.keras.activations import _Activation
 from tensorflow.keras.constraints import Constraint
 from tensorflow.keras.initializers import _Initializer
@@ -90,8 +90,8 @@ class Layer(tf.Module, Generic[_InputT, _OutputT]):
     def non_trainable_weights(self) -> list[Variable]: ...
     @property
     def losses(self) -> list[Tensor]: ...
-    def get_weights(self) -> list[AnyArray]: ...
-    def set_weights(self, weights: Sequence[AnyArray]) -> None: ...
+    def get_weights(self) -> list[_AnyArray]: ...
+    def set_weights(self, weights: Sequence[_AnyArray]) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> Self: ...
diff --git a/stubs/tensorflow/tensorflow/keras/losses.pyi b/stubs/tensorflow/tensorflow/keras/losses.pyi
index e386709c9a55..25aa903bbe65 100644
--- a/stubs/tensorflow/tensorflow/keras/losses.pyi
+++ b/stubs/tensorflow/tensorflow/keras/losses.pyi
@@ -5,7 +5,7 @@ from typing import Any, Final, Literal, TypeVar, overload
 from typing_extensions import Self, TypeAlias, TypeGuard
 
 from tensorflow import Tensor, _TensorCompatible
-from tensorflow._aliases import KerasSerializable
+from tensorflow._aliases import _KerasSerializable
 from tensorflow.keras.metrics import (
     binary_crossentropy as binary_crossentropy,
     categorical_crossentropy as categorical_crossentropy,
@@ -136,7 +136,7 @@ def log_cosh(y_true: _TensorCompatible, y_pred: _TensorCompatible) -> Tensor: ..
 def deserialize(
     name: str | dict[str, Any], custom_objects: dict[str, Any] | None = None, use_legacy_format: bool = False
 ) -> Loss: ...
-def serialize(loss: KerasSerializable, use_legacy_format: bool = False) -> dict[str, Any]: ...
+def serialize(loss: _KerasSerializable, use_legacy_format: bool = False) -> dict[str, Any]: ...
 
 _FuncT = TypeVar("_FuncT", bound=Callable[..., Any])
 
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index a9e3d4799440..7f08d5b08993 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -10,9 +10,18 @@ import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
-from tensorflow.keras.optimizers.legacy import Optimizer
+from tensorflow.keras import _Loss, _Metric
+from tensorflow._aliases import _ContainerGeneric
+
+_BothOptimizer = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
+    _train_counter: tf.Variable
+    _test_counter: tf.Variable
+    optimizer: _BothOptimizer | None
+    loss: tf.keras.losses.Loss | dict[str, tf.keras.losses.Loss]
+    stop_training: bool
+
     def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def __setattr__(self, name: str, value: Any) -> None: ...
@@ -21,13 +30,14 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    # Ideally loss/metrics/output would share the same structure but higher kinded types are not supported.
     def compile(
         self,
-        optimizer: Optimizer | str = "rmsprop",
-        loss: tf.keras.losses.Loss | str | None = None,
-        metrics: list[tf.keras.metrics.Metric | str] | None = None,
-        loss_weights: list[float] | dict[str, float] | None = None,
-        weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
+        optimizer: _BothOptimizer | str = "rmsprop",
+        loss: _ContainerGeneric[_Loss] | None = None,
+        metrics: _ContainerGeneric[_Metric] | None = None,
+        loss_weights: _ContainerGeneric[float] | None = None,
+        weighted_metrics: _ContainerGeneric[_Metric] | None = None,
         run_eagerly: bool | None = None,
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
diff --git a/stubs/tensorflow/tensorflow/keras/optimizers/legacy/__init__.pyi b/stubs/tensorflow/tensorflow/keras/optimizers/legacy/__init__.pyi
index 54ddc40e10fd..ab811879044a 100644
--- a/stubs/tensorflow/tensorflow/keras/optimizers/legacy/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/optimizers/legacy/__init__.pyi
@@ -5,7 +5,7 @@ from typing import Any
 from typing_extensions import Self, TypeAlias
 
 import tensorflow as tf
-from tensorflow._aliases import Gradients
+from tensorflow._aliases import _Gradients
 from tensorflow.keras.optimizers import schedules as schedules
 from tensorflow.python.trackable.base import Trackable
 
@@ -13,9 +13,9 @@ _Initializer: TypeAlias = str | Callable[[], tf.Tensor] | dict[str, Any]
 _Shape: TypeAlias = tf.TensorShape | Iterable[int | None]
 _Dtype: TypeAlias = tf.DType | str | None
 _LearningRate: TypeAlias = float | tf.Tensor | schedules.LearningRateSchedule | Callable[[], float | tf.Tensor]
-_GradientAggregator: TypeAlias = Callable[[list[tuple[Gradients, tf.Variable]]], list[tuple[Gradients, tf.Variable]]] | None
+_GradientAggregator: TypeAlias = Callable[[list[tuple[_Gradients, tf.Variable]]], list[tuple[_Gradients, tf.Variable]]] | None
 _GradientTransformer: TypeAlias = (
-    Iterable[Callable[[list[tuple[Gradients, tf.Variable]]], list[tuple[Gradients, tf.Variable]]]] | None
+    Iterable[Callable[[list[tuple[_Gradients, tf.Variable]]], list[tuple[_Gradients, tf.Variable]]]] | None
 )
 
 # kwargs here and in other optimizers can be given better type after Unpack[TypedDict], PEP 692, is supported.
@@ -53,7 +53,7 @@ class Optimizer(Trackable):
     ) -> tf.Variable: ...
     def apply_gradients(
         self,
-        grads_and_vars: Iterable[tuple[Gradients, tf.Variable]],
+        grads_and_vars: Iterable[tuple[_Gradients, tf.Variable]],
         name: str | None = None,
         experimental_aggregate_gradients: bool = True,
     ) -> tf.Operation | None: ...
@@ -64,7 +64,7 @@ class Optimizer(Trackable):
     def get_config(self) -> dict[str, Any]: ...
     def get_slot(self, var: tf.Variable, slot_name: str) -> tf.Variable: ...
     def get_slot_names(self) -> list[str]: ...
-    def get_gradients(self, loss: tf.Tensor, params: list[tf.Variable]) -> list[Gradients]: ...
+    def get_gradients(self, loss: tf.Tensor, params: list[tf.Variable]) -> list[_Gradients]: ...
     def minimize(
         self,
         loss: tf.Tensor | Callable[[], tf.Tensor],

From 5e75c93ae41052196ba9f7304c7f9c8eac83e009 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:05:18 +0900
Subject: [PATCH 44/77] Fix tf functions.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 3fb6c2bde55c..94bc377ff15c 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -409,12 +409,10 @@ def __getattr__(name: str) -> Incomplete: ...
 def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
-def expand_dims(input: Tensor | tuple[int, ...], axis: int, name: None | str = None) -> Tensor: ...
-def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
+def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
+def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
-def constant(
-    value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = "Const"
-) -> Tensor: ...
+def constant(value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
 def cast(x: Tensor | _TensorCompatible, dtype: DType, name: None | str = None) -> Tensor: ...
 def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...

From 34880648c5d316f92f82333d961b288e1ea3e293 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:15:40 +0900
Subject: [PATCH 45/77] Fix CheckpointOptions's callbacks type.

---
 stubs/tensorflow/tensorflow/train.py | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/train.py

diff --git a/stubs/tensorflow/tensorflow/train.py b/stubs/tensorflow/tensorflow/train.py
new file mode 100644
index 000000000000..dd467d1a0dd4
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/train.py
@@ -0,0 +1,9 @@
+from collections.abc import Callable
+from typing import Any
+
+
+class CheckpointOptions:
+    experimental_io_device: None | str = None
+    experimental_enable_async_checkpoint: bool = False
+    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
+    enable_async: bool = False

From cfd50481b3d86db7eee60f1130028bc60fe639f0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 24 Jan 2024 13:20:35 +0000
Subject: [PATCH 46/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/__init__.pyi     |  4 +++-
 stubs/tensorflow/tensorflow/keras/models.pyi | 19 ++++++++++---------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 94bc377ff15c..bb816c7407e8 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -413,6 +413,8 @@ def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ..
 def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
-def constant(value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = 'Const') -> Tensor: ...
+def constant(
+    value: _TensorCompatible, dtype: None | _DTypeLike = None, shape: None | _ShapeLike = None, name: None | str = "Const"
+) -> Tensor: ...
 def cast(x: Tensor | _TensorCompatible, dtype: DType, name: None | str = None) -> Tensor: ...
 def reshape(tensor: Tensor | _TensorCompatible, shape: _ShapeLike, name: None | str = None) -> Tensor: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7f08d5b08993..9701d87195ed 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -222,12 +222,13 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         expand_nested: bool = False,
         show_trainable: bool = False,
         layer_range: None | list[str] | tuple[str, str] = None,
-    ) -> None: ...
-    @property
-    def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
-    def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
-    def get_weight_paths(self) -> dict[str, tf.Variable]: ...
-    def get_compile_config(self) -> dict[str, Any]: ...
-    def compile_from_config(self, config: dict[str, Any]) -> Self: ...
-    def export(self, filepath: str | Path) -> None: ...
-    def save_spec(self, dynamic_batch: bool = True) -> tuple[tuple[tf.TensorSpec, ...], dict[str, tf.TensorSpec]] | None: ...
+    ): ...
+    def load_weights(
+        self,
+        filepath: str | Path,
+        skip_mismatch: bool = False,
+        by_name: bool = False,
+        options: None | tensorflow.train.CheckpointOptions = None,
+    ): ...
+    def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def call(self, __inputs: _InputT) -> _OutputT: ...

From b7622518a543d67db4728bbfd90dd2c83981723b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:22:04 +0900
Subject: [PATCH 47/77] fix: file name train.py -> train.pyi

---
 stubs/tensorflow/tensorflow/train.py  |  9 ---------
 stubs/tensorflow/tensorflow/train.pyi | 16 ++++------------
 2 files changed, 4 insertions(+), 21 deletions(-)
 delete mode 100644 stubs/tensorflow/tensorflow/train.py

diff --git a/stubs/tensorflow/tensorflow/train.py b/stubs/tensorflow/tensorflow/train.py
deleted file mode 100644
index dd467d1a0dd4..000000000000
--- a/stubs/tensorflow/tensorflow/train.py
+++ /dev/null
@@ -1,9 +0,0 @@
-from collections.abc import Callable
-from typing import Any
-
-
-class CheckpointOptions:
-    experimental_io_device: None | str = None
-    experimental_enable_async_checkpoint: bool = False
-    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
-    enable_async: bool = False
diff --git a/stubs/tensorflow/tensorflow/train.pyi b/stubs/tensorflow/tensorflow/train.pyi
index 8afafdf972f0..49a78be25b98 100644
--- a/stubs/tensorflow/tensorflow/train.pyi
+++ b/stubs/tensorflow/tensorflow/train.pyi
@@ -2,15 +2,7 @@ from collections.abc import Callable
 from typing import Any
 
 class CheckpointOptions:
-    experimental_io_device: None | str
-    experimental_enable_async_checkpoint: bool
-    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]]
-    enable_async: bool
-
-    def __init__(
-        self,
-        experimental_io_device: None | str = None,
-        experimental_enable_async_checkpoint: bool = False,
-        experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None,
-        enable_async: bool = False,
-    ) -> None: ...
+    experimental_io_device: None | str = None
+    experimental_enable_async_checkpoint: bool = False
+    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
+    enable_async: bool = False

From 33f01322900a20f94fc1df984a51111873a330bf Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:30:45 +0900
Subject: [PATCH 48/77] fix: Model MRO error

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 20 ++------------------
 1 file changed, 2 insertions(+), 18 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 9701d87195ed..99ac6c7322a6 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,32 +1,16 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Container, Iterator
 from pathlib import Path
-from typing import Any, Literal
-from typing_extensions import Self
 
 import numpy as np
 import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
-from tensorflow import Variable, _ShapeLike, _TensorCompatible
+from tensorflow import _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
-from tensorflow.keras import _Loss, _Metric
-from tensorflow._aliases import _ContainerGeneric
-
-_BothOptimizer = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
-    _train_counter: tf.Variable
-    _test_counter: tf.Variable
-    optimizer: _BothOptimizer | None
-    loss: tf.keras.losses.Loss | dict[str, tf.keras.losses.Loss]
-    stop_training: bool
-
-    def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
-    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
-    def __setattr__(self, name: str, value: Any) -> None: ...
-    def __reduce__(self) -> Incomplete: ...
-    def __deepcopy__(self, memo: Incomplete) -> Incomplete: ...
+    def __init__(self, *args, **kwargs) -> None: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...

From d67f0b91cc35904dcff25b2e5f75f7bc96e3ff50 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 22:33:33 +0900
Subject: [PATCH 49/77] fix: concat's values type

---
 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 bb816c7407e8..4b65b945109a 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -410,7 +410,7 @@ def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
 def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
-def concat(values: Iterable[Tensor] | Tensor, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
+def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...
 def constant(

From 114f213c097ec72f35e2c9110027a0c149bc3b31 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:42:30 +0900
Subject: [PATCH 50/77] fix: add missing Model methods/properties (wip)

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 261 ++++++++++---------
 1 file changed, 132 insertions(+), 129 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 99ac6c7322a6..ee8e8c19dc61 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,34 +1,40 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Container, Iterator
+from collections.abc import Callable
 from pathlib import Path
+from typing import Any, Iterator, Literal, Self
+
+# import numpy as np
+# import numpy.typing as npt
 
 import numpy as np
 import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
-from tensorflow import _ShapeLike, _TensorCompatible
+from tensorflow import _ShapeLike, _TensorCompatible, Variable
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
+    def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args, **kwargs) -> None: ...
+    def __setattr__(self, name, value) -> None: ...
+    def __reduce__(self) -> Incomplete: ...
+    def __deepcopy__(self, memo) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    # Ideally loss/metrics/output would share the same structure but higher kinded types are not supported.
+    def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def compile(
         self,
-        optimizer: _BothOptimizer | str = "rmsprop",
-        loss: _ContainerGeneric[_Loss] | None = None,
-        metrics: _ContainerGeneric[_Metric] | None = None,
-        loss_weights: _ContainerGeneric[float] | None = None,
-        weighted_metrics: _ContainerGeneric[_Metric] | None = None,
-        run_eagerly: bool | None = None,
-        steps_per_execution: int | Literal["auto"] | None = None,
-        jit_compile: bool | None = None,
-        *,
-        pss_evaluation_shards: int | Literal["auto"] = 0,
-        **kwargs: Any,
-    ) -> None: ...
+        optimizer="rmsprop",
+        loss=None,
+        metrics=None,
+        loss_weights=None,
+        weighted_metrics=None,
+        run_eagerly=None,
+        steps_per_execution=None,
+        jit_compile=None,
+        pss_evaluation_shards=0,
+        **kwargs,
+    ) -> Incomplete: ...
     @property
     def metrics(self) -> list[Incomplete]: ...
     @property
@@ -37,10 +43,10 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def distribute_strategy(self) -> Incomplete: ...  # tf.distribute.Strategy
     @property
     def run_eagerly(self) -> bool: ...
-    # @property
-    # def autotune_steps_per_execution(self) -> Incomplete: ...  # not present at runtime
-    # @property
-    # def steps_per_execution(self) -> int: ...  # Requires a compiled model. # not present at runtime
+    @property
+    def autotune_steps_per_execution(self) -> Incomplete: ...
+    @property
+    def steps_per_execution(self) -> int: ...
     @property
     def jit_compile(self) -> bool: ...
     @property
@@ -54,145 +60,143 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         sample_weight: Incomplete | None = None,
     ) -> tf.Tensor | None: ...
     def compute_metrics(
-        self, x: _TensorCompatible, y: _TensorCompatible, y_pred: _TensorCompatible, sample_weight: Incomplete
+        self,
+        x: _TensorCompatible,
+        y: _TensorCompatible,
+        y_pred: _TensorCompatible,
+        sample_weight: Incomplete,
     ) -> dict[str, float]: ...
     def get_metrics_result(self) -> dict[str, float]: ...
     def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def fit(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        batch_size: int | None = None,
-        epochs: int = 1,
-        verbose: Literal["auto", 0, 1, 2] = "auto",
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        validation_split: float = 0.0,
-        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
-        shuffle: bool = True,
-        class_weight: dict[int, float] | None = None,
-        sample_weight: npt.NDArray[np.float_] | None = None,
-        initial_epoch: int = 0,
-        steps_per_epoch: int | None = None,
-        validation_steps: int | None = None,
-        validation_batch_size: int | None = None,
-        validation_freq: int | Container[int] = 1,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-    ) -> tf.keras.callbacks.History: ...
+        x=None,
+        y=None,
+        batch_size=None,
+        epochs=1,
+        verbose="auto",
+        callbacks=None,
+        validation_split=0.0,
+        validation_data=None,
+        shuffle=True,
+        class_weight=None,
+        sample_weight=None,
+        initial_epoch=0,
+        steps_per_epoch=None,
+        validation_steps=None,
+        validation_batch_size=None,
+        validation_freq=1,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+    ): ...
     def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
     def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def evaluate(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        batch_size: int | None = None,
-        verbose: Literal["auto", 0, 1, 2] = "auto",
-        sample_weight: npt.NDArray[np.float_] | None = None,
-        steps: int | None = None,
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-        return_dict: bool = False,
-        **kwargs: Any,
-    ) -> float | list[float]: ...
+        x=None,
+        y=None,
+        batch_size=None,
+        verbose="auto",
+        sample_weight=None,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        return_dict=False,
+        **kwargs,
+    ): ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x: _TensorCompatible | tf.data.Dataset[Incomplete],
-        batch_size: int | None = None,
-        verbose: Literal["auto", 0, 1, 2] = "auto",
-        steps: int | None = None,
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-    ) -> _OutputT: ...
+        x,
+        batch_size=None,
+        verbose="auto",
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+    ): ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
-        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        sample_weight: npt.NDArray[np.float_] | None = None,
-        class_weight: dict[int, float] | None = None,
-        reset_metrics: bool = True,
-        return_dict: bool = False,
+        x,
+        y=None,
+        sample_weight=None,
+        class_weight=None,
+        reset_metrics=True,
+        return_dict=False,
     ) -> float | list[float]: ...
     def test_on_batch(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
-        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        sample_weight: npt.NDArray[np.float_] | None = None,
-        reset_metrics: bool = True,
-        return_dict: bool = False,
+        x,
+        y=None,
+        sample_weight=None,
+        reset_metrics=True,
+        return_dict=False,
     ) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
-        generator: Iterator[Incomplete],
-        steps_per_epoch: int | None = None,
-        epochs: int = 1,
-        verbose: Literal["auto", 0, 1, 2] = 1,
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
-        validation_steps: int | None = None,
-        validation_freq: int | Container[int] = 1,
-        class_weight: dict[int, float] | None = None,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-        shuffle: bool = True,
-        initial_epoch: int = 0,
-    ) -> tf.keras.callbacks.History: ...
+        generator,
+        steps_per_epoch=None,
+        epochs=1,
+        verbose=1,
+        callbacks=None,
+        validation_data=None,
+        validation_steps=None,
+        validation_freq=1,
+        class_weight=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        shuffle=True,
+        initial_epoch=0,
+    ): ...
     def evaluate_generator(
         self,
-        generator: Iterator[Incomplete],
-        steps: int | None = None,
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-        verbose: Literal["auto", 0, 1, 2] = 0,
-    ) -> float | list[float]: ...
+        generator,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        verbose=0,
+    ): ...
     def predict_generator(
         self,
-        generator: Iterator[Incomplete],
-        steps: int | None = None,
-        callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size: int = 10,
-        workers: int = 1,
-        use_multiprocessing: bool = False,
-        verbose: Literal["auto", 0, 1, 2] = 0,
-    ) -> _OutputT: ...
+        generator,
+        steps=None,
+        callbacks=None,
+        max_queue_size=10,
+        workers=1,
+        use_multiprocessing=False,
+        verbose=0,
+    ): ...
     @property
     def trainable_weights(self) -> list[Variable]: ...
     @property
     def non_trainable_weights(self) -> list[Variable]: ...
-    def get_weights(self) -> Incomplete: ...
-    def save(
-        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
-    ) -> None: ...
+    def get_weights(self): ...
+    def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
     def save_weights(
-        self,
-        filepath: str | Path,
-        overwrite: bool = True,
-        save_format: Literal["tf", "h5"] | None = None,
-        options: tf.train.CheckpointOptions | None = None,
-    ) -> None: ...
+        self, filepath, overwrite=True, save_format=None, options=None
+    ): ...
     def load_weights(
         self,
         filepath: str | Path,
         skip_mismatch: bool = False,
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
-    ) -> None: ...
-    def get_config(self) -> dict[str, Any]: ...
+    ): ...
+    def get_config(self): ...
     @classmethod
-    def from_config(cls, config: dict[str, Any], custom_objects: Incomplete | None = None) -> Self: ...
-    def to_json(self, **kwargs: Any) -> str: ...
-    def to_yaml(self, **kwargs: Any) -> str: ...
+    def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
+    def to_json(self, **kwargs) -> str: ...
+    def to_yaml(self, **kwargs) -> str: ...
     def reset_states(self) -> None: ...
     @property
     def state_updates(self) -> list[Incomplete]: ...
@@ -207,12 +211,11 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         show_trainable: bool = False,
         layer_range: None | list[str] | tuple[str, str] = None,
     ): ...
-    def load_weights(
-        self,
-        filepath: str | Path,
-        skip_mismatch: bool = False,
-        by_name: bool = False,
-        options: None | tensorflow.train.CheckpointOptions = None,
-    ): ...
-    def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    def call(self, __inputs: _InputT) -> _OutputT: ...
+    @property
+    def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
+    def get_layer(self, name=None, index=None) -> Layer[Incomplete, Incomplete]: ...
+    def get_weight_paths(self): ...
+    def get_compile_config(self) -> dict[str, Any]: ...
+    def compile_from_config(self, config: dict[str, Any]) -> Self: ...
+    def export(self, filepath: str | Path) -> None: ...
+    def save_spec(self, dynamic_batch: bool = True) -> tuple[tuple[tf.TensorSpec, ...], dict[str, tf.TensorSpec]] | None: ...

From cd5b7cde3c65f723d61611f2727cac928338f3e6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 24 Jan 2024 14:47:15 +0000
Subject: [PATCH 51/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 50 +++-----------------
 1 file changed, 7 insertions(+), 43 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index ee8e8c19dc61..80bcef91637a 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -5,12 +5,9 @@ from typing import Any, Iterator, Literal, Self
 
 # import numpy as np
 # import numpy.typing as npt
-
-import numpy as np
-import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
-from tensorflow import _ShapeLike, _TensorCompatible, Variable
+from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
@@ -60,11 +57,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         sample_weight: Incomplete | None = None,
     ) -> tf.Tensor | None: ...
     def compute_metrics(
-        self,
-        x: _TensorCompatible,
-        y: _TensorCompatible,
-        y_pred: _TensorCompatible,
-        sample_weight: Incomplete,
+        self, x: _TensorCompatible, y: _TensorCompatible, y_pred: _TensorCompatible, sample_weight: Incomplete
     ) -> dict[str, float]: ...
     def get_metrics_result(self) -> dict[str, float]: ...
     def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
@@ -122,22 +115,9 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ): ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        class_weight=None,
-        reset_metrics=True,
-        return_dict=False,
-    ) -> float | list[float]: ...
-    def test_on_batch(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        reset_metrics=True,
-        return_dict=False,
+        self, x, y=None, sample_weight=None, class_weight=None, reset_metrics=True, return_dict=False
     ) -> float | list[float]: ...
+    def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
@@ -157,24 +137,10 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         initial_epoch=0,
     ): ...
     def evaluate_generator(
-        self,
-        generator,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        verbose=0,
+        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
     ): ...
     def predict_generator(
-        self,
-        generator,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        verbose=0,
+        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
     ): ...
     @property
     def trainable_weights(self) -> list[Variable]: ...
@@ -182,9 +148,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
-    def save_weights(
-        self, filepath, overwrite=True, save_format=None, options=None
-    ): ...
+    def save_weights(self, filepath, overwrite=True, save_format=None, options=None): ...
     def load_weights(
         self,
         filepath: str | Path,

From 6ee7428a5d7c72334a90e270d3a394973bef72ad Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:50:50 +0900
Subject: [PATCH 52/77] fix: use Iterator from collections.abc instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 80bcef91637a..0d642f273cfc 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Any, Iterator, Literal, Self
+from typing import Any, Literal, Self
 
 # import numpy as np
 # import numpy.typing as npt

From 36499f8a1c2a0eb9bf53f6943bec5729737740c5 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 24 Jan 2024 23:51:31 +0900
Subject: [PATCH 53/77] fix: import Self from typing_extensions instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 0d642f273cfc..4b27b8272450 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,7 +1,8 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Any, Literal, Self
+from typing import Any, Literal
+from typing_extensions import Self
 
 # import numpy as np
 # import numpy.typing as npt

From 38f84808504230679136f0443b61a2d6d19b9b6a Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:13:49 +0900
Subject: [PATCH 54/77] feat: add empty Metric class

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index f27206828a99..32b1e3169735 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,6 +1,4 @@
-from typing import Any
-from typing_extensions import Self
-
+from _typeshed import Incomplete
 from tensorflow import Tensor, _TensorCompatible
 from tensorflow.dtypes import DType
 
@@ -8,6 +6,9 @@ class Metric:
     def __init__(self, name: str, dtype: DType) -> None: ...
     def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
 
+class Metric(Incomplete):
+    ...
+
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1
 ) -> Tensor: ...

From 10b432f51ccfca8bcd651df9a0b74b23ac595728 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:17:11 +0900
Subject: [PATCH 55/77] feat: Add type hints for keras Model's compile, fit and
 evaluate

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 93 ++++++++++----------
 1 file changed, 47 insertions(+), 46 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 4b27b8272450..ace5efeec42e 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,15 +1,16 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Container
 from pathlib import Path
 from typing import Any, Literal
 from typing_extensions import Self
 
-# import numpy as np
-# import numpy.typing as npt
+import numpy as np
+import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
+from tensorflow.keras.optimizers.legacy import Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
     def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
@@ -22,17 +23,17 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def compile(
         self,
-        optimizer="rmsprop",
-        loss=None,
-        metrics=None,
-        loss_weights=None,
-        weighted_metrics=None,
-        run_eagerly=None,
-        steps_per_execution=None,
-        jit_compile=None,
-        pss_evaluation_shards=0,
+        optimizer: Optimizer | str = "rmsprop",
+        loss: tf.keras.losses.Loss | str | None = None,
+        metrics: list[tf.keras.metrics.Metric | str] | None = None,
+        loss_weights: list[float] | dict[str, float] | None = None,
+        weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
+        run_eagerly: bool | None = None,
+        steps_per_execution: int |  Literal["auto"] | None= None,
+        jit_compile: bool | None = None,
+        pss_evaluation_shards: int |  Literal["auto"] = 0,
         **kwargs,
-    ) -> Incomplete: ...
+    ) -> None: ...
     @property
     def metrics(self) -> list[Incomplete]: ...
     @property
@@ -64,52 +65,52 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_train_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def fit(
         self,
-        x=None,
-        y=None,
-        batch_size=None,
-        epochs=1,
-        verbose="auto",
-        callbacks=None,
-        validation_split=0.0,
-        validation_data=None,
-        shuffle=True,
-        class_weight=None,
-        sample_weight=None,
-        initial_epoch=0,
-        steps_per_epoch=None,
-        validation_steps=None,
-        validation_batch_size=None,
-        validation_freq=1,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
+        epochs: int = 1,
+        verbose: Literal["auto", 0, 1, 2] = "auto",
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        validation_split: float = 0.0,
+        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
+        shuffle: bool = True,
+        class_weight: dict[int, float] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        initial_epoch: int = 0,
+        steps_per_epoch: int | None = None,
+        validation_steps: int | None = None,
+        validation_batch_size: int | None = None,
+        validation_freq: int | Container[int] = 1,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
     ): ...
     def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
     def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def evaluate(
         self,
-        x=None,
-        y=None,
-        batch_size=None,
-        verbose="auto",
-        sample_weight=None,
-        steps=None,
-        callbacks=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        return_dict=False,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
+        verbose: Literal["auto", 0, 1, 2] = "auto",
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        return_dict: bool = False,
         **kwargs,
     ): ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
         batch_size=None,
-        verbose="auto",
+        verbose: Literal["auto", 0, 1, 2] = "auto",
         steps=None,
-        callbacks=None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
         max_queue_size=10,
         workers=1,
         use_multiprocessing=False,

From 7cae2a00b9ec078dce617a347a9b7942b9a09d3b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:31:08 +0900
Subject: [PATCH 56/77] feat: Add type hints for keras Model's predict,
 train_on_batch, fit_generator, evaluate_generator and predict_generator

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 66 +++++++++++++-------
 1 file changed, 43 insertions(+), 23 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index ace5efeec42e..1e6a3a0f21ca 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -106,43 +106,63 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
-        batch_size=None,
+        x: _TensorCompatible | tf.data.Dataset[Incomplete] | None = None,
+        batch_size: int | None = None,
         verbose: Literal["auto", 0, 1, 2] = "auto",
-        steps=None,
+        steps: int | None = None,
         callbacks: list[tf.keras.callbacks.Callback] | None = None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
     ): ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
-        self, x, y=None, sample_weight=None, class_weight=None, reset_metrics=True, return_dict=False
+        self,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        class_weight: dict[int, float] | None = None,
+        reset_metrics: bool = True,
+        return_dict: bool = False
     ) -> float | list[float]: ...
     def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
-        generator,
-        steps_per_epoch=None,
-        epochs=1,
-        verbose=1,
-        callbacks=None,
-        validation_data=None,
-        validation_steps=None,
-        validation_freq=1,
-        class_weight=None,
-        max_queue_size=10,
-        workers=1,
-        use_multiprocessing=False,
-        shuffle=True,
-        initial_epoch=0,
+        generator: Iterator[Incomplete],
+        steps_per_epoch: int | None = None,
+        epochs: int = 1,
+        verbose: Literal["auto", 0, 1, 2] = 1,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        validation_data: _TensorCompatible | tf.data.Dataset[Any] | None = None,
+        validation_steps: int | None = None,
+        validation_freq: int | Container[int] = 1,
+        class_weight: dict[int, float] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        shuffle: bool = True,
+        initial_epoch: int = 0,
     ): ...
     def evaluate_generator(
-        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
+        self,
+        generator: Iterator[Incomplete],
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        verbose: Literal["auto", 0, 1, 2] = 0,
     ): ...
     def predict_generator(
-        self, generator, steps=None, callbacks=None, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=0
+        self,
+        generator: Iterator[Incomplete],
+        steps: int | None = None,
+        callbacks: list[tf.keras.callbacks.Callback] | None = None,
+        max_queue_size: int = 10,
+        workers: int = 1,
+        use_multiprocessing: bool = False,
+        verbose: Literal["auto", 0, 1, 2] = 0,
     ): ...
     @property
     def trainable_weights(self) -> list[Variable]: ...

From 95ee069c6f06a00f5124aa54053a22658646ceb0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:31:59 +0000
Subject: [PATCH 57/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 4 ++--
 stubs/tensorflow/tensorflow/keras/models.pyi  | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index 32b1e3169735..5959de10bd03 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+
 from tensorflow import Tensor, _TensorCompatible
 from tensorflow.dtypes import DType
 
@@ -6,8 +7,7 @@ class Metric:
     def __init__(self, name: str, dtype: DType) -> None: ...
     def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
 
-class Metric(Incomplete):
-    ...
+class Metric(Incomplete): ...
 
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 1e6a3a0f21ca..9dd8b6bb5be3 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -1,5 +1,5 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Iterator, Container
+from collections.abc import Callable, Container, Iterator
 from pathlib import Path
 from typing import Any, Literal
 from typing_extensions import Self
@@ -29,9 +29,9 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         loss_weights: list[float] | dict[str, float] | None = None,
         weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
         run_eagerly: bool | None = None,
-        steps_per_execution: int |  Literal["auto"] | None= None,
+        steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
-        pss_evaluation_shards: int |  Literal["auto"] = 0,
+        pss_evaluation_shards: int | Literal["auto"] = 0,
         **kwargs,
     ) -> None: ...
     @property
@@ -123,7 +123,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         sample_weight: npt.NDArray[np.float_] | None = None,
         class_weight: dict[int, float] | None = None,
         reset_metrics: bool = True,
-        return_dict: bool = False
+        return_dict: bool = False,
     ) -> float | list[float]: ...
     def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]

From 08a0aa191e7d489f16c07e6de872b60c1a21b65d Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:35:56 +0900
Subject: [PATCH 58/77] feat: Add type hints for keras Model's save,
 save_weights and load_weights

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 9dd8b6bb5be3..84c268ae5aad 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -169,8 +169,20 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
-    def save(self, filepath, overwrite=True, save_format=None, **kwargs): ...
-    def save_weights(self, filepath, overwrite=True, save_format=None, options=None): ...
+    def save(
+        self,
+        filepath: str | Path,
+        overwrite: bool = True,
+        save_format: Literal["keras", "tf", "h5"] | None = None,
+        **kwargs,
+    ): ...
+    def save_weights(
+        self,
+        filepath: str | Path,
+        overwrite: bool = True,
+        save_format: Literal["tf", "h5"] | None = None,
+        options: tf.train.CheckpointOptions | None = None,
+    ): ...
     def load_weights(
         self,
         filepath: str | Path,
@@ -178,7 +190,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
-    def get_config(self): ...
+    def get_config(self) -> None: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
     def to_json(self, **kwargs) -> str: ...

From 0164ed96ce990043e343bf7df89ea3df6e4f0b05 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:37:11 +0900
Subject: [PATCH 59/77] feat: Add type hints for keras Model's get_layer

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 84c268ae5aad..7dbe3a9b858a 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -211,7 +211,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ): ...
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
-    def get_layer(self, name=None, index=None) -> Layer[Incomplete, Incomplete]: ...
+    def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
     def get_weight_paths(self): ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...

From 14ea8de6004c38514332931b667727a9d3510cf2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:38:10 +0000
Subject: [PATCH 60/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7dbe3a9b858a..5b0407548af2 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -170,11 +170,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(
-        self,
-        filepath: str | Path,
-        overwrite: bool = True,
-        save_format: Literal["keras", "tf", "h5"] | None = None,
-        **kwargs,
+        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs
     ): ...
     def save_weights(
         self,

From 2c09fccf3edbbaacaeabd9a6f83a48415bb24a00 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:45:36 +0900
Subject: [PATCH 61/77] fix: add Any type to args and kwargs.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 5b0407548af2..bc14e89b9587 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -13,8 +13,8 @@ from tensorflow.keras.layers import Layer, _InputT, _OutputT
 from tensorflow.keras.optimizers.legacy import Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
-    def __new__(cls, *args, **kwargs) -> Model[_InputT, _OutputT]: ...
-    def __init__(self, *args, **kwargs) -> None: ...
+    def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def __setattr__(self, name, value) -> None: ...
     def __reduce__(self) -> Incomplete: ...
     def __deepcopy__(self, memo) -> Incomplete: ...
@@ -32,7 +32,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
         pss_evaluation_shards: int | Literal["auto"] = 0,
-        **kwargs,
+        **kwargs: Any,
     ) -> None: ...
     @property
     def metrics(self) -> list[Incomplete]: ...
@@ -100,7 +100,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         return_dict: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ): ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
@@ -170,7 +170,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def non_trainable_weights(self) -> list[Variable]: ...
     def get_weights(self): ...
     def save(
-        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs
+        self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
     ): ...
     def save_weights(
         self,
@@ -189,8 +189,8 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def get_config(self) -> None: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
-    def to_json(self, **kwargs) -> str: ...
-    def to_yaml(self, **kwargs) -> str: ...
+    def to_json(self, **kwargs: Any) -> str: ...
+    def to_yaml(self, **kwargs: Any) -> str: ...
     def reset_states(self) -> None: ...
     @property
     def state_updates(self) -> list[Incomplete]: ...

From 2c1603a54cc0d539576d8df113a15ea077e5064d Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:54:01 +0900
Subject: [PATCH 62/77] fix: allow tuple[int, ...] as input to expand_dims

This is not written in t he doc, but is used by jax
---
 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 4b65b945109a..3fb6c2bde55c 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -409,7 +409,7 @@ def __getattr__(name: str) -> Incomplete: ...
 def convert_to_tensor(
     value: _TensorCompatible, dtype: DType | None = None, dtype_hint: DType | None = None, name: str | None = None
 ) -> Tensor: ...
-def expand_dims(input: Tensor, axis: int, name: None | str = None) -> Tensor: ...
+def expand_dims(input: Tensor | tuple[int, ...], axis: int, name: None | str = None) -> Tensor: ...
 def concat(values: _TensorCompatible, axis: int | Tensor, name: None | str = "concat") -> Tensor: ...
 def squeeze(input: Tensor, axis: None | _TensorCompatible = None, name: None | str = None) -> Tensor: ...
 def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor, name: None | str = None) -> Tensor: ...

From 9ad65a2eec63c34d9504ba6b533ebfc61aeb0af2 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:58:40 +0900
Subject: [PATCH 63/77] fix: add empty tensorflow.kerras.callbacks.Callback
 calls

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index e7d68d4c32b7..191a5db61f49 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -1,4 +1,4 @@
 from _typeshed import Incomplete
 
+
 class Callback(Incomplete): ...
-class History(Incomplete): ...

From 3177cd1d3db296be5253d2b0d6b859f6468f09df Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 12:59:26 +0000
Subject: [PATCH 64/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 1 -
 1 file changed, 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 191a5db61f49..7e078b7b4573 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -1,4 +1,3 @@
 from _typeshed import Incomplete
 
-
 class Callback(Incomplete): ...

From eb911e57e8f37dc49a701ab453fd48e2e7138932 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 21:59:35 +0900
Subject: [PATCH 65/77] fix: Model's get_config return type.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index bc14e89b9587..ba3f5839f2b7 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -186,7 +186,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
     ): ...
-    def get_config(self) -> None: ...
+    def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
     def to_json(self, **kwargs: Any) -> str: ...

From 1ba6088319ac68aa9001d722e2d4fbeb467d383b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:13:27 +0900
Subject: [PATCH 66/77] feat: add empty History class

---
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 1 +
 1 file changed, 1 insertion(+)

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 7e078b7b4573..e7d68d4c32b7 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -1,3 +1,4 @@
 from _typeshed import Incomplete
 
 class Callback(Incomplete): ...
+class History(Incomplete): ...

From ff7bf94b525af427e449d212f2735ac3fff6d925 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:14:02 +0900
Subject: [PATCH 67/77] fix: add type hints for test_on_batch, add missing
 return types

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 36 ++++++++++++--------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index ba3f5839f2b7..7bfe829ae27f 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -15,9 +15,9 @@ from tensorflow.keras.optimizers.legacy import Optimizer
 class Model(Layer[_InputT, _OutputT], tf.Module):
     def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
-    def __setattr__(self, name, value) -> None: ...
+    def __setattr__(self, name: str, value: Any) -> None: ...
     def __reduce__(self) -> Incomplete: ...
-    def __deepcopy__(self, memo) -> Incomplete: ...
+    def __deepcopy__(self, memo: Incomplete) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
@@ -84,7 +84,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         max_queue_size: int = 10,
         workers: int = 1,
         use_multiprocessing: bool = False,
-    ): ...
+    ) -> tf.keras.callbacks.History: ...
     def test_step(self, data: _TensorCompatible) -> dict[str, float]: ...
     def make_test_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], dict[str, float]]: ...
     def evaluate(
@@ -101,7 +101,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         use_multiprocessing: bool = False,
         return_dict: bool = False,
         **kwargs: Any,
-    ): ...
+    ) -> float | list[float]: ...
     def predict_step(self, data: _InputT) -> _OutputT: ...
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
@@ -114,7 +114,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         max_queue_size: int = 10,
         workers: int = 1,
         use_multiprocessing: bool = False,
-    ): ...
+    ) -> _OutputT: ...
     def reset_metrics(self) -> None: ...
     def train_on_batch(
         self,
@@ -125,7 +125,15 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         reset_metrics: bool = True,
         return_dict: bool = False,
     ) -> float | list[float]: ...
-    def test_on_batch(self, x, y=None, sample_weight=None, reset_metrics=True, return_dict=False) -> float | list[float]: ...
+    def test_on_batch(
+        self,
+        x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
+        y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
+        sample_weight: npt.NDArray[np.float_] | None = None,
+        class_weight: dict[int, float] | None = None,
+        reset_metrics: bool = True,
+        return_dict: bool = False,
+    ) -> float | list[float]: ...
     def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
     def fit_generator(
         self,
@@ -143,7 +151,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         use_multiprocessing: bool = False,
         shuffle: bool = True,
         initial_epoch: int = 0,
-    ): ...
+    ) -> tf.keras.callbacks.History: ...
     def evaluate_generator(
         self,
         generator: Iterator[Incomplete],
@@ -153,7 +161,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         verbose: Literal["auto", 0, 1, 2] = 0,
-    ): ...
+    ) -> float | list[float]: ...
     def predict_generator(
         self,
         generator: Iterator[Incomplete],
@@ -163,7 +171,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         workers: int = 1,
         use_multiprocessing: bool = False,
         verbose: Literal["auto", 0, 1, 2] = 0,
-    ): ...
+    ) -> _OutputT: ...
     @property
     def trainable_weights(self) -> list[Variable]: ...
     @property
@@ -171,21 +179,21 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def get_weights(self): ...
     def save(
         self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
-    ): ...
+    ) -> None: ...
     def save_weights(
         self,
         filepath: str | Path,
         overwrite: bool = True,
         save_format: Literal["tf", "h5"] | None = None,
         options: tf.train.CheckpointOptions | None = None,
-    ): ...
+    ) -> None: ...
     def load_weights(
         self,
         filepath: str | Path,
         skip_mismatch: bool = False,
         by_name: bool = False,
         options: None | tensorflow.train.CheckpointOptions = None,
-    ): ...
+    ) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
     def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
@@ -204,11 +212,11 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         expand_nested: bool = False,
         show_trainable: bool = False,
         layer_range: None | list[str] | tuple[str, str] = None,
-    ): ...
+    ) -> None: ...
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
     def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
-    def get_weight_paths(self): ...
+    def get_weight_paths(self) -> dict[str,  tf.Variable ]: ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...
     def export(self, filepath: str | Path) -> None: ...

From a7a12dccca4d0d37f23cb07b378b8ec5cfd7a8eb Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 25 Jan 2024 13:14:45 +0000
Subject: [PATCH 68/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7bfe829ae27f..3273ae2b2470 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -216,7 +216,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def layers(self) -> list[Layer[Incomplete, Incomplete]]: ...
     def get_layer(self, name: str | None = None, index: int | None = None) -> Layer[Incomplete, Incomplete]: ...
-    def get_weight_paths(self) -> dict[str,  tf.Variable ]: ...
+    def get_weight_paths(self) -> dict[str, tf.Variable]: ...
     def get_compile_config(self) -> dict[str, Any]: ...
     def compile_from_config(self, config: dict[str, Any]) -> Self: ...
     def export(self, filepath: str | Path) -> None: ...

From 53a68253a23716446c4a98768edb5f4f794a4a49 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:20:27 +0900
Subject: [PATCH 69/77] fix: add missing types

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 3273ae2b2470..f0d0025aa95f 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -176,7 +176,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def trainable_weights(self) -> list[Variable]: ...
     @property
     def non_trainable_weights(self) -> list[Variable]: ...
-    def get_weights(self): ...
+    def get_weights(self) -> Incomplete: ...
     def save(
         self, filepath: str | Path, overwrite: bool = True, save_format: Literal["keras", "tf", "h5"] | None = None, **kwargs: Any
     ) -> None: ...
@@ -196,7 +196,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     ) -> None: ...
     def get_config(self) -> dict[str, Any]: ...
     @classmethod
-    def from_config(cls, config: dict[str, Any], custom_objects=None) -> Self: ...
+    def from_config(cls, config: dict[str, Any], custom_objects: Incomplete | None = None) -> Self: ...
     def to_json(self, **kwargs: Any) -> str: ...
     def to_yaml(self, **kwargs: Any) -> str: ...
     def reset_states(self) -> None: ...

From 4bd2f0b61e520c6f31b35379c78cbd67385c274f Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:29:37 +0900
Subject: [PATCH 70/77] fix: allow call's training to be None

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index f0d0025aa95f..dc825d77fa81 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -20,7 +20,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def __deepcopy__(self, memo: Incomplete) -> Incomplete: ...
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
-    def call(self, inputs: _InputT, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def compile(
         self,
         optimizer: Optimizer | str = "rmsprop",

From fe1599b590d6967a78dd13462ea297da50a3a5e8 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:35:29 +0900
Subject: [PATCH 71/77] fix: predict input type.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index dc825d77fa81..e7a9cc4fb611 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -45,7 +45,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def autotune_steps_per_execution(self) -> Incomplete: ...
     @property
-    def steps_per_execution(self) -> int: ...
+    def steps_per_execution(self) -> int: ...  # Requires a compiled model.
     @property
     def jit_compile(self) -> bool: ...
     @property
@@ -106,7 +106,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def make_predict_function(self, force: bool = False) -> Callable[[tf.data.Iterator[Incomplete]], _OutputT]: ...
     def predict(
         self,
-        x: _TensorCompatible | tf.data.Dataset[Incomplete] | None = None,
+        x: _TensorCompatible | tf.data.Dataset[Incomplete],
         batch_size: int | None = None,
         verbose: Literal["auto", 0, 1, 2] = "auto",
         steps: int | None = None,
@@ -130,7 +130,6 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         x: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete],
         y: _TensorCompatible | dict[str, _TensorCompatible] | tf.data.Dataset[Incomplete] | None = None,
         sample_weight: npt.NDArray[np.float_] | None = None,
-        class_weight: dict[int, float] | None = None,
         reset_metrics: bool = True,
         return_dict: bool = False,
     ) -> float | list[float]: ...

From d9f9669c2f062a36476577cb62a6894daba49af8 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:46:56 +0900
Subject: [PATCH 72/77] fix: fix CheckpointOptions (add init)

---
 stubs/tensorflow/tensorflow/train.pyi | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/train.pyi b/stubs/tensorflow/tensorflow/train.pyi
index 49a78be25b98..8afafdf972f0 100644
--- a/stubs/tensorflow/tensorflow/train.pyi
+++ b/stubs/tensorflow/tensorflow/train.pyi
@@ -2,7 +2,15 @@ from collections.abc import Callable
 from typing import Any
 
 class CheckpointOptions:
-    experimental_io_device: None | str = None
-    experimental_enable_async_checkpoint: bool = False
-    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None
-    enable_async: bool = False
+    experimental_io_device: None | str
+    experimental_enable_async_checkpoint: bool
+    experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]]
+    enable_async: bool
+
+    def __init__(
+        self,
+        experimental_io_device: None | str = None,
+        experimental_enable_async_checkpoint: bool = False,
+        experimental_write_callbacks: None | list[Callable[[str], Any] | Callable[[], Any]] = None,
+        enable_async: bool = False,
+    ) -> None: ...

From 1405abe1ff6919223d86a81f06f2147e09707978 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:47:39 +0900
Subject: [PATCH 73/77] fix: add __init__ and __new__ to Metric

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index 5959de10bd03..e5e474a0e53e 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,4 +1,4 @@
-from _typeshed import Incomplete
+from typing import Any, Self
 
 from tensorflow import Tensor, _TensorCompatible
 from tensorflow.dtypes import DType
@@ -7,8 +7,6 @@ class Metric:
     def __init__(self, name: str, dtype: DType) -> None: ...
     def __new__(cls, *args: Any, **kwargs: Any) -> Self: ...
 
-class Metric(Incomplete): ...
-
 def binary_crossentropy(
     y_true: _TensorCompatible, y_pred: _TensorCompatible, from_logits: bool = False, label_smoothing: float = 0.0, axis: int = -1
 ) -> Tensor: ...

From 054c8cf945205b157944a583d056cbf7c8d0564f Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 22:48:07 +0900
Subject: [PATCH 74/77] fix: comment out properties not present at runtime.

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index e7a9cc4fb611..a9e3d4799440 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -31,6 +31,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         run_eagerly: bool | None = None,
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,
+        *,
         pss_evaluation_shards: int | Literal["auto"] = 0,
         **kwargs: Any,
     ) -> None: ...
@@ -42,10 +43,10 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def distribute_strategy(self) -> Incomplete: ...  # tf.distribute.Strategy
     @property
     def run_eagerly(self) -> bool: ...
-    @property
-    def autotune_steps_per_execution(self) -> Incomplete: ...
-    @property
-    def steps_per_execution(self) -> int: ...  # Requires a compiled model.
+    # @property
+    # def autotune_steps_per_execution(self) -> Incomplete: ...  # not present at runtime
+    # @property
+    # def steps_per_execution(self) -> int: ...  # Requires a compiled model. # not present at runtime
     @property
     def jit_compile(self) -> bool: ...
     @property

From b2a5e7b7f3b2fad4c12eb78e0a0b137453c0d0e7 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Thu, 25 Jan 2024 23:14:56 +0900
Subject: [PATCH 75/77] fix: import Self from typing_extensions instead of
 typing

---
 stubs/tensorflow/tensorflow/keras/metrics.pyi | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/tensorflow/keras/metrics.pyi b/stubs/tensorflow/tensorflow/keras/metrics.pyi
index e5e474a0e53e..f27206828a99 100644
--- a/stubs/tensorflow/tensorflow/keras/metrics.pyi
+++ b/stubs/tensorflow/tensorflow/keras/metrics.pyi
@@ -1,4 +1,5 @@
-from typing import Any, Self
+from typing import Any
+from typing_extensions import Self
 
 from tensorflow import Tensor, _TensorCompatible
 from tensorflow.dtypes import DType

From 04c93eaf087b8efb90dc447f2b5ffa136475381b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Fri, 26 Jan 2024 16:32:11 +0900
Subject: [PATCH 76/77] adding alias types, wip

---
 stubs/tensorflow/tensorflow/keras/models.pyi | 22 ++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index a9e3d4799440..7f08d5b08993 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -10,9 +10,18 @@ import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
-from tensorflow.keras.optimizers.legacy import Optimizer
+from tensorflow.keras import _Loss, _Metric
+from tensorflow._aliases import _ContainerGeneric
+
+_BothOptimizer = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
+    _train_counter: tf.Variable
+    _test_counter: tf.Variable
+    optimizer: _BothOptimizer | None
+    loss: tf.keras.losses.Loss | dict[str, tf.keras.losses.Loss]
+    stop_training: bool
+
     def __new__(cls, *args: Any, **kwargs: Any) -> Model[_InputT, _OutputT]: ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def __setattr__(self, name: str, value: Any) -> None: ...
@@ -21,13 +30,14 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     def build(self, input_shape: _ShapeLike) -> None: ...
     def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ...
     def call(self, inputs: _InputT, training: bool | None = None, mask: _TensorCompatible | None = None) -> _OutputT: ...
+    # Ideally loss/metrics/output would share the same structure but higher kinded types are not supported.
     def compile(
         self,
-        optimizer: Optimizer | str = "rmsprop",
-        loss: tf.keras.losses.Loss | str | None = None,
-        metrics: list[tf.keras.metrics.Metric | str] | None = None,
-        loss_weights: list[float] | dict[str, float] | None = None,
-        weighted_metrics: list[tf.keras.metrics.Metric] | None = None,
+        optimizer: _BothOptimizer | str = "rmsprop",
+        loss: _ContainerGeneric[_Loss] | None = None,
+        metrics: _ContainerGeneric[_Metric] | None = None,
+        loss_weights: _ContainerGeneric[float] | None = None,
+        weighted_metrics: _ContainerGeneric[_Metric] | None = None,
         run_eagerly: bool | None = None,
         steps_per_execution: int | Literal["auto"] | None = None,
         jit_compile: bool | None = None,

From 62f42ffa73bc61dbce95544ffe992b45dda776a0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 27 Jan 2024 08:14:09 +0000
Subject: [PATCH 77/77] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 stubs/tensorflow/tensorflow/_aliases.pyi       |  3 +--
 stubs/tensorflow/tensorflow/keras/__init__.pyi | 17 ++++-------------
 stubs/tensorflow/tensorflow/keras/models.pyi   |  4 ++--
 3 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi
index bf7d4d3ea1db..d2dc5b08fe16 100644
--- a/stubs/tensorflow/tensorflow/_aliases.pyi
+++ b/stubs/tensorflow/tensorflow/_aliases.pyi
@@ -3,9 +3,8 @@
 # equivalent.
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Protocol, TypeVar
+from typing import Any, Iterable, Mapping, Protocol, Sequence, TypeVar
 from typing_extensions import TypeAlias
-from typing import Any, Iterable, Mapping, Sequence, TypeVar
 
 import numpy as np
 import tensorflow as tf
diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 4705b3f2b1ea..0ad0cae06305 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -1,6 +1,8 @@
 from _typeshed import Incomplete
 from typing import Callable
 
+import tensorflow as tf
+from tensorflow._aliases import _TensorCompatible
 from tensorflow.keras import (
     activations as activations,
     constraints as constraints,
@@ -13,20 +15,9 @@ from tensorflow.keras import (
     regularizers as regularizers,
 )
 from tensorflow.keras.models import Model as Model
-import tensorflow as tf
-from tensorflow._aliases import _TensorCompatible
 
 def __getattr__(name: str) -> Incomplete: ...
 
-_Loss = (
-   str
-    | tf.keras.losses.Loss
-    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
-)
+_Loss = str | tf.keras.losses.Loss | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
 
-_Metric = (
-    str
-    | tf.keras.metrics.Metric
-    | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
-    | None
-)
+_Metric = str | tf.keras.metrics.Metric | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor] | None
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 7f08d5b08993..e2f17e8eb3ee 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -9,9 +9,9 @@ import numpy.typing as npt
 import tensorflow
 import tensorflow as tf
 from tensorflow import Variable, _ShapeLike, _TensorCompatible
-from tensorflow.keras.layers import Layer, _InputT, _OutputT
-from tensorflow.keras import _Loss, _Metric
 from tensorflow._aliases import _ContainerGeneric
+from tensorflow.keras import _Loss, _Metric
+from tensorflow.keras.layers import Layer, _InputT, _OutputT
 
 _BothOptimizer = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer