How to set the Keras backend

Keras 3 selects its execution framework before the first keras import in each Python process. Setting the backend at process startup lets one project run the same Keras API on JAX, TensorFlow, or PyTorch without changing model code.

Use KERAS_BACKEND for a project shell, job runner, notebook server, or one-off command because it is explicit and easy to change per launch. The backend package itself must already be installed; keras does not install JAX, TensorFlow, or PyTorch as part of the backend switch.

For a user-wide default, Keras also reads ~/.keras/keras.json when no environment override is present. Once keras has been imported, changing the environment or config file does not move that running process to another backend, so restart notebooks, shells, workers, and long-running servers after changing the selection.

1. Choose the backend name that matches the installed framework package.

   Use jax, tensorflow, or torch for normal Keras 3 training and inference workflows. Install the matching backend framework before launching Python.

2. Set the backend for a one-off command before importing keras.

   $ KERAS_BACKEND=jax python3 -c 'import keras;print(keras.config.backend())'
   jax

   The inline environment assignment applies only to that command in POSIX-style shells.

3. Export the backend in the shell that launches the project.

   $ export KERAS_BACKEND=jax

4. Verify a fresh Python process sees the exported backend.

   $ python3 -c 'import keras;print(keras.config.backend())'
   jax

5. Set the backend inside a Python entrypoint when the script or notebook owns startup.

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
    
   print(keras.config.backend())

   Put the environment assignment before importing keras, keras_hub, keras_cv, or any project module that imports Keras indirectly.

6. Set a user-wide default only when every new Keras process should use the same backend.

   {
       "backend": "jax",
       "floatx": "float32",
       "epsilon": 1e-07,
       "image_data_format": "channels_last"
   }

   Save the JSON in ~/.keras/keras.json. A KERAS_BACKEND environment variable still overrides this file for the process launch.

7. Clear the shell override before testing the config-file default.

   $ unset KERAS_BACKEND

8. Verify a new Python process reads the default backend.

   $ python3 -c 'import keras;print(keras.config.backend())'
   jax

   If the output still shows a different backend, restart the Python process and check whether another startup file, notebook kernel, or parent shell exported KERAS_BACKEND.

   Related: How to fix Keras backend selection after import