How do you choose the Keras 3 backend, and when must that choice happen?
answer
- one environment variable, three frameworks
- import order is the whole trick
- a JSON config file holds the persistent default
- archives cross backends, custom code may not
basics
~20 sSet the KERAS_BACKEND environment variable to tensorflow, jax or torch before the first import of keras, or set the backend field in the ~/.keras/keras.json config file. The choice is fixed for the process once keras is imported; keras.backend.backend() reports it.
solid answer
~40 sKeras 3 runs on TensorFlow, JAX or PyTorch, selected per process. In code that means `os.environ["KERAS_BACKEND"] = "jax"` **before** `import keras`, because the backend is resolved at import time and wired into the ops dispatch; setting it afterwards does nothing. The persistent alternative is the `backend` field in `~/.keras/keras.json`, and the environment variable overrides that file. Verify with `keras.backend.backend()`, which returns the active backend name. What travels across backends is the saved model: a `.keras` archive is config JSON plus plain weight arrays, so a model of built-in layers trained under JAX loads under PyTorch. What does not travel is custom code written against one framework's API - which is what `keras.ops` exists to avoid.
code
python · 5 linesimport os
os.environ["KERAS_BACKEND"] = "torch"
import keras
print(keras.backend.backend()) # torchgo deeper
Remember that Keras 3 runs on TensorFlow, JAX or PyTorch and that you pick one with the KERAS_BACKEND environment variable set before importing keras.
Explain why the ordering matters: the backend is resolved at import and bound into the dispatch layer, so a later change to the variable has no effect until the process restarts.
Discuss what actually crosses the seam - portable archives, custom code that may not - and log keras.backend.backend() at startup so environment drift is visible rather than inferred.
Own the choice as an infrastructure decision: pick the backend that matches your data pipelines, deployment tooling and hiring, then enforce keras.ops in shared code so the decision stays reversible.
## Selecting a backend Keras 3 is a multi-backend API: the same Dense, the same fit(), executing on TensorFlow, JAX or PyTorch. Selection is per process, by one of two mechanisms. The environment variable, used in scripts and containers: ``` import os os.environ["KERAS_BACKEND"] = "jax" import keras # must come after ``` Or the config file `~/.keras/keras.json`, whose `backend` field sets a machine-wide default (alongside `floatx` and `epsilon`). The environment variable takes precedence over the file. `keras.backend.backend()` returns the active name, and it is worth logging at startup - most "why is this slower here" and "why did this op fail" confusion resolves the moment you see which backend is live. ## Why the ordering rule is strict At import, Keras resolves the backend once and binds its numeric dispatch to that framework. Mutating `os.environ` afterwards changes a string nobody reads again. In a notebook this bites constantly: cell 1 imports keras, cell 5 sets KERAS_BACKEND, and everything keeps running on the old backend until the kernel restarts. Set it in the very first cell, before any import that might pull Keras in transitively. The backend framework must of course be installed; Keras does not vendor them. ## What each backend buys you The API is the same, the surrounding ecosystem is not. TensorFlow brings tf.data input pipelines and TensorFlow's deployment tooling. JAX brings its compilation and sharding story. PyTorch brings the PyTorch ecosystem - and under that backend a Keras layer is also a torch.nn.Module, so a Keras model can be dropped into a PyTorch training loop or wrapped by PyTorch tooling. Teams usually choose the backend that matches their surrounding infrastructure rather than the one that trains marginally faster. ## What crosses the seam Saved models cross it. The `.keras` archive holds config.json (class names and arguments) and model.weights.h5 (plain numeric arrays) - nothing backend-specific. Train under JAX, load under PyTorch, and a model of built-in layers reconstructs and runs. Custom code is where portability is won or lost. A layer whose call() uses keras.ops is backend-agnostic; a layer that calls one framework's tensor API directly pins the saved model to that backend, even though the archive itself looks portable. Data pipelines are the other seam. Keras 3 accepts several input types, but a pipeline built on one framework's dataset object is a dependency on that framework in your training script, even if the model is portable. ## The Keras 2 mental model to unlearn In Keras 2, Keras was TensorFlow's front end and the tensors flowing through it were tf.Tensor. Writing Keras 3 code that assumes that is a portability bug waiting to be discovered. The tell in an answer is someone reaching for a TensorFlow symbol to describe a Keras 3 operation instead of keras.ops.
- You set KERAS_BACKEND in a notebook cell but the backend does not change. What happened?Keras was already imported, so the backend had been resolved and bound at import time; the variable is read once. Restart the kernel and set it in the first cell before any import, including imports of your own modules that pull Keras in transitively. keras.backend.backend() confirms which one you actually got.
- Can a .keras model trained under TensorFlow be loaded under the PyTorch backend?Yes, if the model is built from Keras layers and any custom code is backend-agnostic. The archive stores a config of class names and arguments plus weights as plain arrays, so Keras rebuilds the layers under whichever backend is active. It fails only when a custom layer calls one framework's ops directly.
saying these in an interview costs you the question
- Setting KERAS_BACKEND after importing keras
- Thinking Keras 3 is still TensorFlow-only
- Expecting Keras to install the backend framework
- Believing saved weights are backend-specific