optimizers.MovingAverage not working with tf2.11

[Ianni]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04
• TensorFlow version and how it was installed (source or binary): 2.11 pip wheel
• TensorFlow-Addons version and how it was installed (source or binary): 0.19 pip
• Python version: 3.9
• Is GPU used? (yes/no): yes

Describe the bug

Using moving average of optimizers is no longer working and results in

TypeError: optimizer is not an object of tf.keras.optimizers.Optimizer or tf.keras.optimizers.legacy.Optimizer (if you have tf version >= 2.11.0).

Code to reproduce the issue

Compare e.g. the example notebook from the documentation:
https://www.tensorflow.org/addons/tutorials/average_optimizers_callback

Note: if you try it in collab, you have to install tf2.11 first.

Other info / logs

See also https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/MovingAverage

[bhack]

/cc @chenmoneygithub

[bhack]

/cc @chenmoneygithub

[bobertlo]

I got this also. I encountered the same error on SWA, then found a collar notebook using MovingAverage before SWA and that has the same error in the collar runtime as I was getting on my laptop.

[bertsky]

As the error message already states: with TF >= 2.11, you have to wrap an tf.keras.optimizers.legacy.Optimizer rather than tf.keras.optimizers.Optimizer (because the latter is now what used to be tf.keras.optimizers.experimental.Optimizer).