From 145c6d58c431c8c6b6235a83fb6b749e3a45e9ec Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 28 Jan 2024 20:24:51 +0900
Subject: [PATCH 01/24] begin adding models.Model

---
 .../tensorflow/tensorflow/keras/__init__.pyi  |   8 +
 stubs/tensorflow/tensorflow/keras/models.pyi  | 233 ++++++++++++++++++
 2 files changed, 241 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/keras/models.pyi

diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 09b1b4737528..a2bd13bff629 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -1,5 +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,
@@ -7,8 +10,13 @@ from tensorflow.keras import (
     layers as layers,
     losses as losses,
     metrics as metrics,
+    models as models,
     optimizers as optimizers,
     regularizers as regularizers,
 )
+from tensorflow.keras.models import Model as Model
+
+_Loss = str | tf.keras.losses.Loss | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor]
+_Metric = str | tf.keras.metrics.Metric | Callable[[_TensorCompatible, _TensorCompatible], tf._Tensor] | None
 
 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..d2680acaec26
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -0,0 +1,233 @@
+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._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
+
+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 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: _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: ...
+    @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: ...  # 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
+    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: _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: ...
+    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]: ...
+    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: ...
+    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,
+    ) -> 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,
+    ) -> 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: ...
+    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]: ...
+    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: ...
+    @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 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: Incomplete | None = None) -> Self: ...
+    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]: ...
+    @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,
+    ) -> 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: ...

From 538a2547b67a8822018a1c86a32c91e8bf8c1384 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 31 Jan 2024 22:20:23 +0900
Subject: [PATCH 02/24] add Model's call, __new__ and build to allowlist
 (similar to Layer)

---
 stubs/tensorflow/@tests/stubtest_allowlist.txt | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
index 5700e283375f..21f699b70148 100644
--- a/stubs/tensorflow/@tests/stubtest_allowlist.txt
+++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -60,10 +60,13 @@ tensorflow.keras.layers.*.__init__
 tensorflow.keras.layers.*.call
 tensorflow.keras.regularizers.Regularizer.__call__
 tensorflow.keras.constraints.Constraint.__call__
+# The same as above applies to Model.call.
+tensorflow.keras.models.*.call
 
-# Layer class does good deal of __new__ magic and actually returns one of two different internal
+# Layer/Model class does good deal of __new__ magic and actually returns one of two different internal
 # types depending on tensorflow execution mode. This feels like implementation internal.
 tensorflow.keras.layers.Layer.__new__
+tensorflow.keras.model.Model.__new__
 
 # build/compute_output_shape are marked positional only in stubs
 # as argument name is inconsistent across layer's and looks like
@@ -71,6 +74,8 @@ tensorflow.keras.layers.Layer.__new__
 # disagreements.
 tensorflow.keras.layers.*.build
 tensorflow.keras.layers.*.compute_output_shape
+# The same as above applies to Model.build.
+tensorflow.keras.models.*.build
 
 # pb2.pyi generated by mypy-protobuf diverge with runtime in many ways. These stubs
 # are mainly tested in mypy-protobuf.

From ca07f575445380c18a4fea62e3a04978ba57d65c Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sat, 17 Feb 2024 15:18:09 +0900
Subject: [PATCH 03/24] uncomment code.

---
 stubs/tensorflow/tensorflow/distribute/__init__.pyi |  5 +++++
 stubs/tensorflow/tensorflow/keras/models.pyi        | 12 ++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)
 create mode 100644 stubs/tensorflow/tensorflow/distribute/__init__.pyi

diff --git a/stubs/tensorflow/tensorflow/distribute/__init__.pyi b/stubs/tensorflow/tensorflow/distribute/__init__.pyi
new file mode 100644
index 000000000000..87bbb0425c99
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/distribute/__init__.pyi
@@ -0,0 +1,5 @@
+from _typeshed import Incomplete
+
+Strategy = Incomplete
+
+def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 94ceb32f628c..02ce20ea7727 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -51,13 +51,13 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     @property
     def metrics_names(self) -> list[str]: ...
     @property
-    def distribute_strategy(self) -> Incomplete: ...  # tf.distribute.Strategy
+    def distribute_strategy(self) -> 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) -> bool: ...
+    @property
+    def steps_per_execution(self) -> int | None: ...  # Returns None for a non-compiled model.
     @property
     def jit_compile(self) -> bool: ...
     @property
@@ -145,7 +145,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
         reset_metrics: bool = True,
         return_dict: bool = False,
     ) -> float | list[float]: ...
-    def predict_on_batch(self, x: Iterator[_InputT]) -> Incomplete: ...  # npt.NDArray[_OutputT]
+    def predict_on_batch(self, x: Iterator[_InputT]) -> npt.NDArray[Incomplete]: ...
     def fit_generator(
         self,
         generator: Iterator[Incomplete],

From c11817d0dff580b673d5f260faff9a12460143c9 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sat, 17 Feb 2024 15:32:02 +0900
Subject: [PATCH 04/24] fix: fix model in callbacks

---
 stubs/tensorflow/tensorflow/keras/__init__.pyi  |  1 +
 stubs/tensorflow/tensorflow/keras/callbacks.pyi | 10 +++++-----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi
index 3a43597e1c6d..4aa06e9b9e67 100644
--- a/stubs/tensorflow/tensorflow/keras/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi
@@ -2,6 +2,7 @@ from _typeshed import Incomplete
 
 from tensorflow.keras import (
     activations as activations,
+    callbacks as callbacks,
     constraints as constraints,
     initializers as initializers,
     layers as layers,
diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 1e632e9c0fa5..1499ea61b4a8 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -12,9 +12,9 @@ from tensorflow.train import CheckpointOptions
 _Logs: TypeAlias = Mapping[str, Any] | None | Any
 
 class Callback:
-    model: Model  # Model[Any, object]
+    model: Model[Any, Any]
     params: dict[str, Any]
-    def set_model(self, model: Model) -> None: ...
+    def set_model(self, model: Model[Any, Any]) -> None: ...
     def set_params(self, params: dict[str, Any]) -> None: ...
     def on_batch_begin(self, batch: int, logs: _Logs = None) -> None: ...
     def on_batch_end(self, batch: int, logs: _Logs = None) -> None: ...
@@ -35,7 +35,7 @@ class Callback:
 
 # A CallbackList has exact same api as a callback, but does not actually subclass it.
 class CallbackList:
-    model: Model
+    model: Model[Any, Any]
     params: dict[str, Any]
     def __init__(
         self,
@@ -43,10 +43,10 @@ class CallbackList:
         add_history: bool = False,
         add_progbar: bool = False,
         # model: Model[Any, object] | None = None,
-        model: Model | None = None,
+        model: Model[Any, Any] | None = None,
         **params: Any,
     ) -> None: ...
-    def set_model(self, model: Model) -> None: ...
+    def set_model(self, model: Model[Any, Any]) -> None: ...
     def set_params(self, params: dict[str, Any]) -> None: ...
     def on_batch_begin(self, batch: int, logs: _Logs | None = None) -> None: ...
     def on_batch_end(self, batch: int, logs: _Logs | None = None) -> None: ...

From efdda57588b804fa11d34f249a32b69fec54786e Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sat, 17 Feb 2024 15:54:26 +0900
Subject: [PATCH 05/24] Remove optimizer alias

---
 stubs/tensorflow/tensorflow/keras/models.pyi              | 6 +++---
 stubs/tensorflow/tensorflow/keras/optimizers/__init__.pyi | 2 ++
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 02ce20ea7727..3b9b3730579e 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -11,15 +11,15 @@ import tensorflow as tf
 from tensorflow import Variable
 from tensorflow._aliases import ContainerGeneric, ShapeLike, TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
+from tensorflow.keras.optimizers import Optimizer
 
-_BothOptimizer: TypeAlias = tf.optimizers.Optimizer | tf.optimizers.experimental.Optimizer
 _Loss: TypeAlias = str | tf.keras.losses.Loss | Callable[[TensorCompatible, TensorCompatible], tf._Tensor]
 _Metric: TypeAlias = str | tf.keras.metrics.Metric | Callable[[TensorCompatible, TensorCompatible], tf._Tensor] | None
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
     _train_counter: tf.Variable
     _test_counter: tf.Variable
-    optimizer: _BothOptimizer | None
+    optimizer: Optimizer | None
     loss: tf.keras.losses.Loss | dict[str, tf.keras.losses.Loss]
     stop_training: bool
 
@@ -34,7 +34,7 @@ class Model(Layer[_InputT, _OutputT], tf.Module):
     # Ideally loss/metrics/output would share the same structure but higher kinded types are not supported.
     def compile(
         self,
-        optimizer: _BothOptimizer | str = "rmsprop",
+        optimizer: Optimizer | str = "rmsprop",
         loss: ContainerGeneric[_Loss] | None = None,
         metrics: ContainerGeneric[_Metric] | None = None,
         loss_weights: ContainerGeneric[float] | None = None,
diff --git a/stubs/tensorflow/tensorflow/keras/optimizers/__init__.pyi b/stubs/tensorflow/tensorflow/keras/optimizers/__init__.pyi
index 7d5346e613f2..a1ccef2f9a4f 100644
--- a/stubs/tensorflow/tensorflow/keras/optimizers/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/keras/optimizers/__init__.pyi
@@ -2,4 +2,6 @@ from _typeshed import Incomplete
 
 from tensorflow.keras.optimizers import legacy as legacy, schedules as schedules
 
+Optimizer = Incomplete
+
 def __getattr__(name: str) -> Incomplete: ...

From c58c9aa2ffda4ed008d657e8e830f02a72968f49 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sat, 17 Feb 2024 15:59:29 +0900
Subject: [PATCH 06/24] fix: Tensor's name

---
 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 3b9b3730579e..46f8cd3ba9ef 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -13,8 +13,8 @@ from tensorflow._aliases import ContainerGeneric, ShapeLike, TensorCompatible
 from tensorflow.keras.layers import Layer, _InputT, _OutputT
 from tensorflow.keras.optimizers import Optimizer
 
-_Loss: TypeAlias = str | tf.keras.losses.Loss | Callable[[TensorCompatible, TensorCompatible], tf._Tensor]
-_Metric: TypeAlias = str | tf.keras.metrics.Metric | Callable[[TensorCompatible, TensorCompatible], tf._Tensor] | None
+_Loss: TypeAlias = str | tf.keras.losses.Loss | Callable[[TensorCompatible, TensorCompatible], tf.Tensor]
+_Metric: TypeAlias = str | tf.keras.metrics.Metric | Callable[[TensorCompatible, TensorCompatible], tf.Tensor] | None
 
 class Model(Layer[_InputT, _OutputT], tf.Module):
     _train_counter: tf.Variable

From ca259ae54b0e01dea8369a138586816b74bcc2cf Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 10:28:03 +0900
Subject: [PATCH 07/24] add pre-existing stubs.

Stubs from: https://github.com/hmc-cs-mdrissi/tensorflow_stubs/blob/main/stubs/tensorflow/saved_model/__init__.pyi
---
 .../tensorflow/saved_model/__init__.pyi       | 72 +++++++++++++++++++
 .../tensorflow/saved_model/experimental.pyi   |  0
 2 files changed, 72 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/saved_model/__init__.pyi
 create mode 100644 stubs/tensorflow/tensorflow/saved_model/experimental.pyi

diff --git a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
new file mode 100644
index 000000000000..05caa8f855fd
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
@@ -0,0 +1,72 @@
+from typing import Any, Generic, Mapping, Sequence, TypeVar
+from typing_extensions import ParamSpec
+
+import tensorflow as tf
+from tensorflow.compat.v1.saved_model.signature_constants import *
+from tensorflow.python.training.tracking.autotrackable import AutoTrackable
+from tensorflow.saved_model.experimental import VariablePolicy
+from tensorflow.types.experimental import ConcreteFunction, GenericFunction
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R", covariant=True)
+
+class SaveOptions:
+    __slots__ = (
+        "namespace_whitelist",
+        "save_debug_info",
+        "function_aliases",
+        "experimental_io_device",
+        "experimental_variable_policy",
+        "experimental_custom_gradients",
+    )
+    namespace_whitelist: list[str]
+    save_debug_info: bool
+    function_aliases: dict[str, tf.types.experimental.GenericFunction[..., object]]
+    experimental_io_device: str
+    experimental_variable_policy: VariablePolicy
+    experimental_custom_gradients: bool
+    def __init__(
+        self,
+        namespace_whitelist: list[str] | None = None,
+        save_debug_info: bool = False,
+        function_aliases: dict[str, tf.types.experimental.GenericFunction[..., object]] | None = None,
+        experimental_io_device: str | None = None,
+        experimental_variable_policy: str | VariablePolicy | None = None,
+        experimental_custom_gradients: bool = True,
+    ): ...
+
+class LoadOptions:
+    def __init__(
+        self,
+        allow_partial_checkpoint: bool = False,
+        experimental_io_device: str | None = None,
+        experimental_skip_checkpoint: bool = False,
+    ) -> None: ...
+
+class _LoadedAttributes(Generic[_P, _R]):
+    signatures: Mapping[str, ConcreteFunction[_P, _R]]
+
+class _LoadedModel(AutoTrackable, _LoadedAttributes[_P, _R]):
+    variables: list[tf.Variable]
+    trainable_variables: list[tf.Variable]
+    # TF1 model artifact specific
+    graph: tf.Graph
+
+def load(
+    export_dir: str, tags: str | Sequence[str] | None = None, options: LoadOptions | None = None
+) -> _LoadedModel[..., Any]: ...
+
+_TF_Function = ConcreteFunction[..., object] | GenericFunction[..., object]
+
+def save(
+    obj: tf.Module,
+    export_dir: str,
+    signatures: _TF_Function | Mapping[str, _TF_Function] | None = None,
+    options: SaveOptions | None = None,
+) -> None: ...
+
+SERVING = "serve"
+TRAINING = "train"
+EVAL = "eval"
+GPU = "gpu"
+TPU = "tpu"
diff --git a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
new file mode 100644
index 000000000000..e69de29bb2d1

From 17865894a9e8ec0cf54ae180bd2b2524e52505b2 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 12:39:12 +0900
Subject: [PATCH 08/24] add saved_model.experimental

---
 stubs/tensorflow/tensorflow/_aliases.pyi      |  2 +
 .../tensorflow/python/trackable/ressource.pyi |  5 +++
 .../tensorflow/saved_model/experimental.pyi   | 37 +++++++++++++++++++
 3 files changed, 44 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/python/trackable/ressource.pyi

diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi
index 842c5e5acd47..a067ea931cb9 100644
--- a/stubs/tensorflow/tensorflow/_aliases.pyi
+++ b/stubs/tensorflow/tensorflow/_aliases.pyi
@@ -27,6 +27,8 @@ class KerasSerializable2(Protocol):
 
 KerasSerializable: TypeAlias = KerasSerializable1 | KerasSerializable2
 
+TensorValue: TypeAlias = tf.Tensor  # Alias for a 0D Tensor
+Integer: TypeAlias = TensorValue | int | IntArray | np.number[Any]  # Here IntArray are assumed to be 0D.
 Slice: TypeAlias = int | slice | None
 FloatDataSequence: TypeAlias = Sequence[float] | Sequence[FloatDataSequence]
 IntDataSequence: TypeAlias = Sequence[int] | Sequence[IntDataSequence]
diff --git a/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi b/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi
new file mode 100644
index 000000000000..58dffc42de13
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi
@@ -0,0 +1,5 @@
+from tensorflow.python.trackable.base import Trackable
+
+# Internal type that is commonly used as a base class
+# and some public apis the signature needs it.
+class CapturableResource(Trackable): ...
diff --git a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
index e69de29bb2d1..218e13230427 100644
--- a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
@@ -0,0 +1,37 @@
+from _typeshed import Incomplete
+from enum import Enum
+from typing import Self
+
+import tensorflow as tf
+from tensorflow._aliases import Integer, TensorValue
+from tensorflow.python.trackable.ressource import CapturableResource
+
+class Fingerprint:
+    saved_model_checksum: TensorValue | None
+    graph_def_program_hash: TensorValue | None = None
+    signature_def_hash: TensorValue | None = None
+    saved_object_graph_hash: TensorValue | None = None
+    checkpoint_hash: TensorValue | None = None
+    version: TensorValue | None = None
+    # In practice it seems like any type is accepted, but that might cause issues later on.
+    def __init__(
+        self,
+        saved_model_checksum: Integer | None = None,
+        graph_def_program_hash: Integer | None = None,
+        signature_def_hash: Integer | None = None,
+        saved_object_graph_hash: Integer | None = None,
+        checkpoint_hash: Integer | None = None,
+        version: Integer | None = None,
+    ) -> None: ...
+    @classmethod
+    def from_proto(cls, proto: Incomplete) -> Self: ...
+    def singleprint(self) -> str: ...
+
+class TrackableResource(CapturableResource):
+    resource_handle: tf.Tensor
+    def __init__(self, device: str = "") -> None: ...
+
+class VariablePolicy(Enum):
+    EXPAND_DISTRIBUTED_VARIABLES = "expand_distributed_variables"
+    NONE = None  # noqa: Y026
+    SAVE_VARIABLE_DEVICES = "save_variable_devices"

From 986df8c140408ff98961b5a0e02a173f8b75e3dc Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 12:42:25 +0900
Subject: [PATCH 09/24] add read_fingerprint

---
 stubs/tensorflow/tensorflow/saved_model/experimental.pyi | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
index 218e13230427..49e1ad033d72 100644
--- a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
@@ -35,3 +35,5 @@ class VariablePolicy(Enum):
     EXPAND_DISTRIBUTED_VARIABLES = "expand_distributed_variables"
     NONE = None  # noqa: Y026
     SAVE_VARIABLE_DEVICES = "save_variable_devices"
+
+def read_fingerprint(export_dir: str) -> Fingerprint: ...

From 4282a1fffa034157961f8ede3edcd12ea00507dc Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:02:52 +0900
Subject: [PATCH 10/24] Add empty func_graph

---
 stubs/tensorflow/tensorflow/python/framework/func_graph.pyi | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/python/framework/func_graph.pyi

diff --git a/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi b/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi
new file mode 100644
index 000000000000..0a41363bf82b
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi
@@ -0,0 +1,3 @@
+from _typeshed import Incomplete
+
+FuncGraph = Incomplete

From bac52fd4a2956a095bed648934fa7659ec779b9b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:03:35 +0900
Subject: [PATCH 11/24] Add empty autotrackable

---
 .../tensorflow/python/training/tracking/autotrackable.pyi      | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/python/training/tracking/autotrackable.pyi

diff --git a/stubs/tensorflow/tensorflow/python/training/tracking/autotrackable.pyi b/stubs/tensorflow/tensorflow/python/training/tracking/autotrackable.pyi
new file mode 100644
index 000000000000..96f4743c44d3
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/python/training/tracking/autotrackable.pyi
@@ -0,0 +1,3 @@
+from _typeshed import Incomplete
+
+AutoTrackable = Incomplete

From 057e9e741fc0f0ce5b07534de82fb6066dd04cc0 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:05:17 +0900
Subject: [PATCH 12/24] add types.experimental from pre-existing stubs

Taken from: https://github.com/hmc-cs-mdrissi/tensorflow_stubs/blob/main/stubs/tensorflow/types/experimental.pyi
---
 .../tensorflow/types/experimental.pyi         | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)
 create mode 100644 stubs/tensorflow/tensorflow/types/experimental.pyi

diff --git a/stubs/tensorflow/tensorflow/types/experimental.pyi b/stubs/tensorflow/tensorflow/types/experimental.pyi
new file mode 100644
index 000000000000..52f41403cc78
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/types/experimental.pyi
@@ -0,0 +1,32 @@
+from typing import Any, Generic, TypeVar, overload
+from typing_extensions import ParamSpec
+
+import tensorflow as tf
+from tensorflow._aliases import ContainerGeneric
+from tensorflow.python.framework.func_graph import FuncGraph
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R", covariant=True)
+
+class Callable(Generic[_P, _R]):
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: ...
+
+class GenericFunction(Callable[_P, _R]):
+    @overload
+    def get_concrete_function(self, *args: _P.args, **kwargs: _P.kwargs) -> ConcreteFunction[_P, _R]: ...
+    @overload
+    def get_concrete_function(  # type: ignore
+        self, *args: ContainerGeneric[tf.TypeSpec[Any]], **kwargs: ContainerGeneric[tf.TypeSpec[Any]]
+    ) -> ConcreteFunction[_P, _R]: ...
+
+class ConcreteFunction(Callable[_P, _R]):
+    @property
+    def structured_input_signature(self) -> tuple[tuple[Any, ...], dict[str, tf.TypeSpec[Any]]]: ...
+    @property
+    def structured_outputs(self) -> dict[str, tf.Tensor | tf.SparseTensor | tf.RaggedTensor | tf.TypeSpec[Any]]: ...
+    @property
+    def graph(self) -> FuncGraph: ...
+    @property
+    def inputs(self) -> list[tf.Tensor]: ...
+    @property
+    def outputs(self) -> list[tf.Tensor]: ...

From 748537dde087a65ee0d36cd89d86d490e871a2b9 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:07:46 +0900
Subject: [PATCH 13/24] complete
 stubs/tensorflow/tensorflow/saved_model/__init__.pyi

---
 .../tensorflow/saved_model/__init__.pyi       | 72 +++++++++++++++----
 1 file changed, 57 insertions(+), 15 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
index 05caa8f855fd..d8044fa5d9a1 100644
--- a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
@@ -1,8 +1,9 @@
-from typing import Any, Generic, Mapping, Sequence, TypeVar
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any, Generic, Literal, TypeAlias, TypeVar
 from typing_extensions import ParamSpec
 
 import tensorflow as tf
-from tensorflow.compat.v1.saved_model.signature_constants import *
 from tensorflow.python.training.tracking.autotrackable import AutoTrackable
 from tensorflow.saved_model.experimental import VariablePolicy
 from tensorflow.types.experimental import ConcreteFunction, GenericFunction
@@ -10,6 +11,28 @@ from tensorflow.types.experimental import ConcreteFunction, GenericFunction
 _P = ParamSpec("_P")
 _R = TypeVar("_R", covariant=True)
 
+class Asset:
+    asset_path: tf.Tensor
+    def __init__(self, path: str | Path | tf.Tensor) -> None: ...
+
+class LoadOptions:
+    allow_partial_checkpoint: bool
+    experimental_io_device: str | None
+    experimental_skip_checkpoint: bool
+    experimental_variable_policy: VariablePolicy | None
+    experimental_load_function_aliases: bool
+
+    def __init__(
+        self,
+        allow_partial_checkpoint: bool = False,
+        experimental_io_device: str | None = None,
+        experimental_skip_checkpoint: bool = False,
+        experimental_variable_policy: (
+            VariablePolicy | Literal["expand_distributed_variables", "save_variable_devices"] | None
+        ) = None,
+        experimental_load_function_aliases: bool = False,
+    ) -> None: ...
+
 class SaveOptions:
     __slots__ = (
         "namespace_whitelist",
@@ -18,6 +41,8 @@ class SaveOptions:
         "experimental_io_device",
         "experimental_variable_policy",
         "experimental_custom_gradients",
+        "experimental_image_format",
+        "experimental_skip_saver",
     )
     namespace_whitelist: list[str]
     save_debug_info: bool
@@ -25,6 +50,8 @@ class SaveOptions:
     experimental_io_device: str
     experimental_variable_policy: VariablePolicy
     experimental_custom_gradients: bool
+    experimental_image_format: bool
+    experimental_skip_saver: bool
     def __init__(
         self,
         namespace_whitelist: list[str] | None = None,
@@ -33,15 +60,11 @@ class SaveOptions:
         experimental_io_device: str | None = None,
         experimental_variable_policy: str | VariablePolicy | None = None,
         experimental_custom_gradients: bool = True,
+        experimental_image_format: bool = False,
+        experimental_skip_saver: bool = False,
     ): ...
 
-class LoadOptions:
-    def __init__(
-        self,
-        allow_partial_checkpoint: bool = False,
-        experimental_io_device: str | None = None,
-        experimental_skip_checkpoint: bool = False,
-    ) -> None: ...
+def contains_saved_model(export_dir: str | Path) -> bool: ...
 
 class _LoadedAttributes(Generic[_P, _R]):
     signatures: Mapping[str, ConcreteFunction[_P, _R]]
@@ -56,7 +79,7 @@ def load(
     export_dir: str, tags: str | Sequence[str] | None = None, options: LoadOptions | None = None
 ) -> _LoadedModel[..., Any]: ...
 
-_TF_Function = ConcreteFunction[..., object] | GenericFunction[..., object]
+_TF_Function: TypeAlias = ConcreteFunction[..., object] | GenericFunction[..., object]
 
 def save(
     obj: tf.Module,
@@ -65,8 +88,27 @@ def save(
     options: SaveOptions | None = None,
 ) -> None: ...
 
-SERVING = "serve"
-TRAINING = "train"
-EVAL = "eval"
-GPU = "gpu"
-TPU = "tpu"
+ASSETS_DIRECTORY: str = "assets"
+ASSETS_KEY: str = "saved_model_assets"
+CLASSIFY_INPUTS: str = "inputs"
+CLASSIFY_METHOD_NAME: str = "tensorflow/serving/classify"
+CLASSIFY_OUTPUT_CLASSES: str = "classes"
+CLASSIFY_OUTPUT_SCORES: str = "scores"
+DEBUG_DIRECTORY: str = "debug"
+DEBUG_INFO_FILENAME_PB: str = "saved_model_debug_info.pb"
+DEFAULT_SERVING_SIGNATURE_DEF_KEY: str = "serving_default"
+GPU: str = "gpu"
+PREDICT_INPUTS: str = "inputs"
+PREDICT_METHOD_NAME: str = "tensorflow/serving/predict"
+PREDICT_OUTPUTS: str = "outputs"
+REGRESS_INPUTS: str = "inputs"
+REGRESS_METHOD_NAME: str = "tensorflow/serving/regress"
+REGRESS_OUTPUTS: str = "outputs"
+SAVED_MODEL_FILENAME_PB: str = "saved_model.pb"
+SAVED_MODEL_FILENAME_PBTXT: str = "saved_model.pbtxt"
+SAVED_MODEL_SCHEMA_VERSION: int = 1
+SERVING: str = "serve"
+TPU: str = "tpu"
+TRAINING: str = "train"
+VARIABLES_DIRECTORY: str = "variables"
+VARIABLES_FILENAME: str = "variables"

From 42895b4d23306da96420b620c8e9f8f9c985107a Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:13:04 +0900
Subject: [PATCH 14/24] fix import

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

diff --git a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
index 49e1ad033d72..5934c1849495 100644
--- a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
@@ -1,6 +1,6 @@
 from _typeshed import Incomplete
 from enum import Enum
-from typing import Self
+from typing_extensions import Self
 
 import tensorflow as tf
 from tensorflow._aliases import Integer, TensorValue

From 3cbd71a5981966427dd1d69c764a3e3fe45be980 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:36:15 +0900
Subject: [PATCH 15/24] fix some typing issues (runtime comparison)

---
 .../python/framework/func_graph.pyi           |  3 ---
 .../tensorflow/python/trackable/ressource.pyi |  4 +++-
 .../tensorflow/saved_model/__init__.pyi       |  7 +++---
 .../tensorflow/saved_model/experimental.pyi   |  3 ++-
 .../tensorflow/types/experimental.pyi         | 24 +++++++------------
 5 files changed, 18 insertions(+), 23 deletions(-)
 delete mode 100644 stubs/tensorflow/tensorflow/python/framework/func_graph.pyi

diff --git a/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi b/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi
deleted file mode 100644
index 0a41363bf82b..000000000000
--- a/stubs/tensorflow/tensorflow/python/framework/func_graph.pyi
+++ /dev/null
@@ -1,3 +0,0 @@
-from _typeshed import Incomplete
-
-FuncGraph = Incomplete
diff --git a/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi b/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi
index 58dffc42de13..617bc82544ba 100644
--- a/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi
+++ b/stubs/tensorflow/tensorflow/python/trackable/ressource.pyi
@@ -1,5 +1,7 @@
 from tensorflow.python.trackable.base import Trackable
 
+class _ResourceMetaclass(type): ...
+
 # Internal type that is commonly used as a base class
 # and some public apis the signature needs it.
-class CapturableResource(Trackable): ...
+class CapturableResource(Trackable, metaclass=_ResourceMetaclass): ...
diff --git a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
index d8044fa5d9a1..7e290477c059 100644
--- a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
@@ -1,7 +1,7 @@
 from collections.abc import Mapping, Sequence
 from pathlib import Path
-from typing import Any, Generic, Literal, TypeAlias, TypeVar
-from typing_extensions import ParamSpec
+from typing import Any, Generic, Literal, TypeVar
+from typing_extensions import ParamSpec, TypeAlias
 
 import tensorflow as tf
 from tensorflow.python.training.tracking.autotrackable import AutoTrackable
@@ -12,7 +12,8 @@ _P = ParamSpec("_P")
 _R = TypeVar("_R", covariant=True)
 
 class Asset:
-    asset_path: tf.Tensor
+    @property
+    def asset_path(self) -> tf.Tensor: ...
     def __init__(self, path: str | Path | tf.Tensor) -> None: ...
 
 class LoadOptions:
diff --git a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
index 5934c1849495..578b248ae64c 100644
--- a/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/experimental.pyi
@@ -28,7 +28,8 @@ class Fingerprint:
     def singleprint(self) -> str: ...
 
 class TrackableResource(CapturableResource):
-    resource_handle: tf.Tensor
+    @property
+    def resource_handle(self) -> tf.Tensor: ...
     def __init__(self, device: str = "") -> None: ...
 
 class VariablePolicy(Enum):
diff --git a/stubs/tensorflow/tensorflow/types/experimental.pyi b/stubs/tensorflow/tensorflow/types/experimental.pyi
index 52f41403cc78..e3cba33bbb73 100644
--- a/stubs/tensorflow/tensorflow/types/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/types/experimental.pyi
@@ -1,32 +1,26 @@
+import abc
+from _typeshed import Incomplete
 from typing import Any, Generic, TypeVar, overload
 from typing_extensions import ParamSpec
 
 import tensorflow as tf
 from tensorflow._aliases import ContainerGeneric
-from tensorflow.python.framework.func_graph import FuncGraph
 
 _P = ParamSpec("_P")
 _R = TypeVar("_R", covariant=True)
 
-class Callable(Generic[_P, _R]):
+class Callable(Generic[_P, _R], metaclass=abc.ABCMeta):
     def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: ...
 
-class GenericFunction(Callable[_P, _R]):
+class ConcreteFunction(Callable[_P, _R], metaclass=abc.ABCMeta):
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: ...
+
+class GenericFunction(Callable[_P, _R], metaclass=abc.ABCMeta):
     @overload
     def get_concrete_function(self, *args: _P.args, **kwargs: _P.kwargs) -> ConcreteFunction[_P, _R]: ...
     @overload
-    def get_concrete_function(  # type: ignore
+    def get_concrete_function(
         self, *args: ContainerGeneric[tf.TypeSpec[Any]], **kwargs: ContainerGeneric[tf.TypeSpec[Any]]
     ) -> ConcreteFunction[_P, _R]: ...
 
-class ConcreteFunction(Callable[_P, _R]):
-    @property
-    def structured_input_signature(self) -> tuple[tuple[Any, ...], dict[str, tf.TypeSpec[Any]]]: ...
-    @property
-    def structured_outputs(self) -> dict[str, tf.Tensor | tf.SparseTensor | tf.RaggedTensor | tf.TypeSpec[Any]]: ...
-    @property
-    def graph(self) -> FuncGraph: ...
-    @property
-    def inputs(self) -> list[tf.Tensor]: ...
-    @property
-    def outputs(self) -> list[tf.Tensor]: ...
+def __getattr__(name: str) -> Incomplete: ...

From 2e1dbf835c3681fbc9d526bd082a9605a248104c Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:48:11 +0900
Subject: [PATCH 16/24] add missing return type

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

diff --git a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
index 7e290477c059..5344689cf6db 100644
--- a/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/saved_model/__init__.pyi
@@ -63,7 +63,7 @@ class SaveOptions:
         experimental_custom_gradients: bool = True,
         experimental_image_format: bool = False,
         experimental_skip_saver: bool = False,
-    ): ...
+    ) -> None: ...
 
 def contains_saved_model(export_dir: str | Path) -> bool: ...
 

From ed5c54c00326006124a4dfebe61fee998d60bb8e Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:58:23 +0900
Subject: [PATCH 17/24] Make GenericFunction's get_concrete_function abstract.

---
 stubs/tensorflow/tensorflow/types/experimental.pyi | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/stubs/tensorflow/tensorflow/types/experimental.pyi b/stubs/tensorflow/tensorflow/types/experimental.pyi
index e3cba33bbb73..be0531318fd1 100644
--- a/stubs/tensorflow/tensorflow/types/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/types/experimental.pyi
@@ -17,8 +17,10 @@ class ConcreteFunction(Callable[_P, _R], metaclass=abc.ABCMeta):
 
 class GenericFunction(Callable[_P, _R], metaclass=abc.ABCMeta):
     @overload
+    @abc.abstractmethod
     def get_concrete_function(self, *args: _P.args, **kwargs: _P.kwargs) -> ConcreteFunction[_P, _R]: ...
     @overload
+    @abc.abstractmethod
     def get_concrete_function(
         self, *args: ContainerGeneric[tf.TypeSpec[Any]], **kwargs: ContainerGeneric[tf.TypeSpec[Any]]
     ) -> ConcreteFunction[_P, _R]: ...

From bf1a41d5d420e5ff55d56f394c1276c19230d191 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 13:59:45 +0900
Subject: [PATCH 18/24] add missing GenericFunction method

---
 stubs/tensorflow/tensorflow/types/experimental.pyi | 1 +
 1 file changed, 1 insertion(+)

diff --git a/stubs/tensorflow/tensorflow/types/experimental.pyi b/stubs/tensorflow/tensorflow/types/experimental.pyi
index be0531318fd1..8e1c02495242 100644
--- a/stubs/tensorflow/tensorflow/types/experimental.pyi
+++ b/stubs/tensorflow/tensorflow/types/experimental.pyi
@@ -24,5 +24,6 @@ class GenericFunction(Callable[_P, _R], metaclass=abc.ABCMeta):
     def get_concrete_function(
         self, *args: ContainerGeneric[tf.TypeSpec[Any]], **kwargs: ContainerGeneric[tf.TypeSpec[Any]]
     ) -> ConcreteFunction[_P, _R]: ...
+    def experimental_get_compiler_ir(self, *args: Incomplete, **kwargs: Incomplete) -> Incomplete: ...
 
 def __getattr__(name: str) -> Incomplete: ...

From a151415a0adbdf82cbbf3fe5d53543bc608ea58b Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 14:08:42 +0900
Subject: [PATCH 19/24] fix SaveOptions import

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

diff --git a/stubs/tensorflow/tensorflow/keras/callbacks.pyi b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
index 1499ea61b4a8..a0dbb676c8f5 100644
--- a/stubs/tensorflow/tensorflow/keras/callbacks.pyi
+++ b/stubs/tensorflow/tensorflow/keras/callbacks.pyi
@@ -7,6 +7,7 @@ import tensorflow as tf
 from requests.api import _HeadersMapping
 from tensorflow.keras import Model
 from tensorflow.keras.optimizers.schedules import LearningRateSchedule
+from tensorflow.saved_model import SaveOptions
 from tensorflow.train import CheckpointOptions
 
 _Logs: TypeAlias = Mapping[str, Any] | None | Any
@@ -115,7 +116,7 @@ class LearningRateScheduler(Callback):
 class ModelCheckpoint(Callback):
     monitor_op: Any
     filepath: str
-    _options: CheckpointOptions | tf.saved_model.SaveOptions | None
+    _options: CheckpointOptions | SaveOptions | None
     def __init__(
         self,
         filepath: str,
@@ -125,7 +126,7 @@ class ModelCheckpoint(Callback):
         save_weights_only: bool = False,
         mode: Literal["auto", "min", "max"] = "auto",
         save_freq: str | int = "epoch",
-        options: CheckpointOptions | tf.saved_model.SaveOptions | None = None,
+        options: CheckpointOptions | SaveOptions | None = None,
         initial_value_threshold: float | None = None,
     ) -> None: ...
     def _save_model(self, epoch: int, batch: int | None, logs: _Logs) -> None: ...

From 77f7a8db30239b8f9e230adef3314491d21b88ef Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 14:28:28 +0900
Subject: [PATCH 20/24] make pss_evaluation_shards not keyword-only

---
 stubs/tensorflow/tensorflow/distribute/__init__.pyi          | 2 +-
 stubs/tensorflow/tensorflow/keras/models.pyi                 | 1 -
 .../tensorflow/python/distribute/distribute_lib.pyi          | 5 +++++
 3 files changed, 6 insertions(+), 2 deletions(-)
 create mode 100644 stubs/tensorflow/tensorflow/python/distribute/distribute_lib.pyi

diff --git a/stubs/tensorflow/tensorflow/distribute/__init__.pyi b/stubs/tensorflow/tensorflow/distribute/__init__.pyi
index 87bbb0425c99..e129db61aab8 100644
--- a/stubs/tensorflow/tensorflow/distribute/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/distribute/__init__.pyi
@@ -1,5 +1,5 @@
 from _typeshed import Incomplete
 
-Strategy = Incomplete
+from tensorflow.python.distribute.distribute_lib import Strategy as Strategy
 
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/keras/models.pyi b/stubs/tensorflow/tensorflow/keras/models.pyi
index 46f8cd3ba9ef..2d00100ecd40 100644
--- a/stubs/tensorflow/tensorflow/keras/models.pyi
+++ b/stubs/tensorflow/tensorflow/keras/models.pyi
@@ -42,7 +42,6 @@ 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: ...
diff --git a/stubs/tensorflow/tensorflow/python/distribute/distribute_lib.pyi b/stubs/tensorflow/tensorflow/python/distribute/distribute_lib.pyi
new file mode 100644
index 000000000000..87bbb0425c99
--- /dev/null
+++ b/stubs/tensorflow/tensorflow/python/distribute/distribute_lib.pyi
@@ -0,0 +1,5 @@
+from _typeshed import Incomplete
+
+Strategy = Incomplete
+
+def __getattr__(name: str) -> Incomplete: ...

From b61885ffab7546c9a1cef1a07aea158fb8635118 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Sun, 18 Feb 2024 14:40:22 +0900
Subject: [PATCH 21/24] try to fix jax issues.

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

diff --git a/stubs/tensorflow/tensorflow/train/__init__.pyi b/stubs/tensorflow/tensorflow/train/__init__.pyi
index 591998230129..4e4d76bd3553 100644
--- a/stubs/tensorflow/tensorflow/train/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/train/__init__.pyi
@@ -16,6 +16,7 @@ from tensorflow.core.example.feature_pb2 import (
 from tensorflow.core.protobuf.cluster_pb2 import ClusterDef as ClusterDef
 from tensorflow.core.protobuf.tensorflow_server_pb2 import ServerDef as ServerDef
 from tensorflow.python.trackable.base import Trackable
+from tensorflow.python.training.tracking.autotrackable import AutoTrackable
 
 class CheckpointOptions:
     experimental_io_device: None | str
@@ -44,7 +45,7 @@ class _CheckpointLoadStatus:
     def assert_nontrivial_match(self) -> Self: ...
     def expect_partial(self) -> Self: ...
 
-class Checkpoint:
+class Checkpoint(AutoTrackable):
     def __init__(self, root: Trackable | None = None, **kwargs: Trackable) -> None: ...
     def read(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...
     def restore(self, save_path: str, options: CheckpointOptions | None = None) -> _CheckpointLoadStatus: ...

From b11dd657bca64204184758f8ba0582c0697831c2 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 28 Feb 2024 22:09:44 +0900
Subject: [PATCH 22/24] Add tensorflow.python to allowlist

---
 stubs/tensorflow/@tests/stubtest_allowlist.txt | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
index f0d9080d7bcc..53bd8f57aa91 100644
--- a/stubs/tensorflow/@tests/stubtest_allowlist.txt
+++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -108,3 +108,9 @@ tensorflow.train.FloatList.*
 tensorflow.train.Int64List.*
 tensorflow.train.ClusterDef.*
 tensorflow.train.ServerDef.*
+
+# The python module cannot be accessed directly, so to stubtest it appears that it is not present at runtime.
+# However it can be accessed by doing:
+# from tensorflow import python
+# python.X
+tensorflow.python.*

From 058bed349c7ac65299d01c2d933bafdf6681fa37 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 28 Feb 2024 22:15:15 +0900
Subject: [PATCH 23/24] Add tensorflow.distribute.Strategy to allowlist

---
 stubs/tensorflow/@tests/stubtest_allowlist.txt | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
index 53bd8f57aa91..b423ad276c19 100644
--- a/stubs/tensorflow/@tests/stubtest_allowlist.txt
+++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -114,3 +114,6 @@ tensorflow.train.ServerDef.*
 # from tensorflow import python
 # python.X
 tensorflow.python.*
+
+# The modules below are re-exported from tensorflow.python, and they therefore appear missing to stubtest.
+tensorflow.distribute.Strategy

From 73e38d301399afe5478f8945f16735691de9d253 Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Wed, 28 Feb 2024 22:19:39 +0900
Subject: [PATCH 24/24] remove unused allowlist entries

---
 stubs/tensorflow/@tests/stubtest_allowlist.txt | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
index b423ad276c19..a8f544a7b260 100644
--- a/stubs/tensorflow/@tests/stubtest_allowlist.txt
+++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -68,13 +68,10 @@ tensorflow.keras.layers.*.__init__
 tensorflow.keras.layers.*.call
 tensorflow.keras.regularizers.Regularizer.__call__
 tensorflow.keras.constraints.Constraint.__call__
-# The same as above applies to Model.call.
-tensorflow.keras.models.*.call
 
 # Layer/Model class does good deal of __new__ magic and actually returns one of two different internal
 # types depending on tensorflow execution mode. This feels like implementation internal.
 tensorflow.keras.layers.Layer.__new__
-tensorflow.keras.model.Model.__new__
 
 # build/compute_output_shape are marked positional only in stubs
 # as argument name is inconsistent across layer's and looks like
@@ -82,8 +79,6 @@ tensorflow.keras.model.Model.__new__
 # disagreements.
 tensorflow.keras.layers.*.build
 tensorflow.keras.layers.*.compute_output_shape
-# The same as above applies to Model.build.
-tensorflow.keras.models.*.build
 
 # pb2.pyi generated by mypy-protobuf diverge with runtime in many ways. These stubs
 # are mainly tested in mypy-protobuf.