Cannot checkpoint Optimizer as its not a trackable object.

[sampathweb]

Check pointing optimizer fails in TF 2.16 with Keras 3. We have may tensorflow.org tutorials that do similar things - https://www.tensorflow.org/tutorials/generative/pix2pix/

This code snippet works in TF 2.15

import numpy as np
import tensorflow as tf

print(tf.__version__)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(1)
])
optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
x, y = np.random.random((10, 2)), np.random.random((10, 1))
model.compile(optimizer, "mse")
model.fit(x, y)

checkpoint = tf.train.Checkpoint(model) # Works
checkpoint = tf.train.Checkpoint(model, optimizer=optimizer)  # Fails in TF 2.16 when we checkpoint optimizer

/usr/local/lib/python3.10/dist-packages/tensorflow/python/checkpoint/checkpoint.py in _assert_trackable(obj, name)
   1571   if not isinstance(
   1572       obj, (base.Trackable, def_function.Function)):
-> 1573     raise ValueError(
   1574         f"`Checkpoint` was expecting {name} to be a trackable object (an "
   1575         f"object derived from `Trackable`), got {obj}. If you believe this "

ValueError: `Checkpoint` was expecting optimizer to be a trackable object (an object derived from `Trackable`), got <keras.src.optimizers.adam.Adam object at 0x7bd306cedfc0>. If you believe this object should be trackable (i.e. it is part of the TensorFlow Python API and manages state), please open an issue.

[fchollet]

Thanks for the report. Will take up this one.

[fchollet]

Done.