Bert training with TPU does not work on Keras Core

[martin-gorner]

Repro notebook: https://www.kaggle.com/code/alexia/kerasnlp-starter-notebook-contradictory-dearwatson

This notebook is configured to use keras_nlp and standard tf.keras. It works perfectly.

If reconfigured to use keras_nlp with Keras Core the model stops working (failing version here):

• it displays non-sensical accuracies > 1 during training
• eval accuracy the same as if the mode was doing random predictions
• .predict returns predictions of the wrong shape (shape=(n,) instead of (n,3)
• therefore the final np.argmax(predictions, axis=1) fails.
• there is nan in the predictions

The behavior does not change when pip installing form GitHub master (latest version) rather than PyPi (latest published package) for both Keras Core and KerasNLP

The config I used to make the model fail:

!pip install keras-core
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import keras_core as keras

[sampathweb]

Yes. I confirm the problem of accuracies > 1 in TPU env with keras-core backend. The notebook works fine in GPU env with keras-core backend.