TopolearnChoosing and switching backends is the first decision in any Keras 3 project and affects speed, deployment and ecosystem.
Keras 3 is a rewrite of Keras that is no longer tied to TensorFlow. Models, layers, losses and optimizers are written once against Keras's own API and executed by a backend: "jax", "tensorflow" or "torch" (plus "numpy" and "openvino" for inference only).
Select the backend before Keras is imported, with the KERAS_BACKEND environment variable or the "backend" field of ~/.keras/keras.json. Once imported, the backend cannot change in that process. keras.config.backend() tells you which one is active.
Why it matters:
- JAX is often the fastest for training thanks to XLA compilation, and works well on TPUs.
- TensorFlow gives SavedModel export, TF Serving and LiteRT.
- PyTorch lets Keras models interoperate with PyTorch code: a Keras model is also a
torch.nn.Module.
Tensors you get back are native to the backend (a jax.Array, tf.Tensor or torch.Tensor), but Keras functions accept NumPy arrays everywhere. Install with pip install keras plus the backend framework you want.
In TensorFlow 2.16+, tf.keras also points to Keras 3. Old Keras 2 code that relies on TensorFlow internals may need updates; the tf_keras package keeps legacy Keras 2 available.
Example
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow" / "torch"; set BEFORE import
import keras
print(keras.__version__) # 3.x
print(keras.config.backend()) # jax
x = keras.ops.convert_to_tensor([1.0, 2.0])
print(type(x).__module__) # a native jax array type under JAX
print(keras.ops.convert_to_numpy(keras.ops.sum(x * 2))) # 6.0Remember
- One codebase, three backends: JAX, TensorFlow, PyTorch
- Set
KERAS_BACKENDbefore importing keras - Returned tensors are backend-native; NumPy inputs work everywhere
Watch out for
- Setting
KERAS_BACKENDafterimport keras, which has no effect.