How to disable JIT compilation in Keras

Keras can run a model through just-in-time compilation so the active backend can optimize training or prediction, but compiler support does not cover every backend operation. Disabling JIT is useful when the model logic is valid in normal execution but the traceback names XLA, JIT, or unsupported compiled operations.

The setting lives on the compiled model. Passing jit_compile=False to model.compile() leaves layers, weights, optimizer, loss, metrics, and backend choice unchanged while Keras builds the training or prediction function without the JIT compiler.

Use the change as a targeted troubleshooting step, not as a catch-all for every training error. Shape mismatches, missing backend packages, invalid losses, and NaN values need their own fix; the JIT workaround is for failures that occur while Keras is compiling a function for the selected backend.

1. Create disable_jit_compile.py with TensorFlow backend selection and a model that exposes a JIT failure.

   disable_jit_compile.py

   import os
   import re
    
   os.environ["KERAS_BACKEND"] = "tensorflow"
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import keras
   import numpy as np
   import tensorflow as tf
    
    
   class StringRoundTripModel(keras.Model):
       def call(self, inputs):
           text = tf.strings.as_string(inputs)
           return tf.strings.to_number(text)
    
    
   def summarize_error(error):
       message = str(error)
       if "AsString" in message and "XLA_CPU_JIT" in message:
           return f"{type(error).__name__}: unsupported operation on XLA_CPU_JIT (AsString)"
       detail = "XLA compile failure detected"
       for line in message.splitlines():
           if "Detected unsupported operations" in line:
               detail = re.sub(r"__inference_[A-Za-z0-9_]+", "__inference_model_call", line.strip())
               break
       return f"{type(error).__name__}: {detail}"
    
    
   x_predict = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype="float32")

   Set KERAS_BACKEND before importing keras. This sample uses TensorFlow because the string operation raises a clear XLA unsupported-operation error.
   Related: How to set the Keras backend

2. Add the failing compiled run below the input array.

   broken_model = StringRoundTripModel()
   broken_model.compile(jit_compile=True)
    
   print(f"backend: {keras.backend.backend()}")
   print(f"broken jit_compile: {broken_model.jit_compile}")
    
   try:
       broken_model.predict(x_predict, verbose=0)
   except Exception as error:
       print(f"jit failure: {summarize_error(error)}")
   else:
       raise RuntimeError("jit_compile=True unexpectedly succeeded")

   Project code can fail because it sets jit_compile=True or because the backend chooses a compiled path automatically. The fix is the same compile-time override.

3. Add the corrected run with jit_compile=False below the failure check.

   fixed_model = StringRoundTripModel()
   fixed_model.compile(jit_compile=False)
   prediction = fixed_model.predict(x_predict, verbose=0)
    
   print(f"fixed jit_compile: {fixed_model.jit_compile}")
   print(f"prediction shape: {prediction.shape}")
   print(f"finite output: {np.isfinite(prediction).all()}")
   print(f"first row: {prediction[0].tolist()}")

   Use the same argument in the project model.compile() call that runs before the failing fit(), evaluate(), or predict() call. Do not add a second compile call only for this flag.

4. Run the script and confirm that the model predicts after JIT is disabled.

   $ python3 disable_jit_compile.py
   backend: tensorflow
   broken jit_compile: True
   jit failure: InvalidArgumentError: unsupported operation on XLA_CPU_JIT (AsString)
   fixed jit_compile: False
   prediction shape: (2, 3)
   finite output: True
   first row: [1.0, 2.0, 3.0]