keras/src/trainers/trainer.py/model_supports_jit
def model_supports_jit(model):
    # XLA not supported with TF on MacOS GPU
    if platform.system() == "Darwin" and "arm" in platform.processor().lower():
        if backend.backend() == "tensorflow":
            from keras.src.utils.module_utils import tensorflow as tf

            if tf.config.list_physical_devices("GPU"):
                return False
    # XLA not supported by some layers
    if all(x.supports_jit for x in model._flatten_layers()):
        return True
    return False