AdamW optimiser crashes on tf-macos v2.12.0

Hi,

Apologies if this is not the right place for bug reports in tf-macos, I am happy to move this issue elsewhere if so!

The problem occurs in tensorflow-macos v2.12.0 when attempting to call model.compile() with the AdamW optimiser. A warning is thrown, telling us that there is a known slowdown when using v2.11+ optimizers, and the backend attempts to fall back to a legacy version. However, AdamW does not exist in legacy versions, which eventually leads to an "unknown optimizer" error.

See below for a MWE. I am running on MacBook M1 Pro.

Cheers, Ste

##
##  Imports
##
 
import sys
 
import tensorflow as tf
 
from tensorflow.keras.models     import Model
from tensorflow.keras.layers     import Input, Dense
from tensorflow.keras.optimizers import AdamW
 
##
##  Report versions
##
 
print(f"Python version is: {sys.version}")
##  -->  Python version is: 3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
 
print(f"TF version is: {tf.__version__}")
##  -->  TF version is: 2.12.0
 
print(f"Keras version is: {tf.keras.__version__}")
##  -->  Keras version is: 2.12.0
 
 
##
##  Create a very simple model
##
 
x_in  = Input(1)
x     = Dense(10)(x_in)
model = Model(x_in, x)
 
##
##  Compile model with AdamW optimizer
##
model.compile(optimizer=AdamW(learning_rate=1e-3, weight_decay=1e-2))

which outputs:

It is worth noting that tf.keras.optimizers.legacy.AdamW does not exist, and so we cannot simply use this.