Keras 3 Multi-Backend: Write Once, Run on TensorFlow, JAX, or PyTorch

The portability problem in one scenario

You inherit a forecasting model built with tf.keras. The research team wants to train on JAX to use TPU pods, while the serving team deploys PyTorch models on GPU instances. Maintaining separate copies of the same architecture for TensorFlow, JAX, and PyTorch means every bug fix and feature change has to be repeated three times.

Keras 3 introduces a backend abstraction layer: the same model definition, training loop, and .keras serialization can run on TensorFlow, JAX, or PyTorch. The useful takeaway is that portability is a contract, not a default. Code that reaches into backend-specific tensor APIs will break when you switch backends. Code that stays inside keras.ops and the unified Keras API can move between backends without conversion tools.

What Keras 3 actually abstracts

Keras 3 dispatches core operations—layers, models, optimizers, losses, metrics, and callbacks—to backend-specific kernels. A custom layer should use keras.ops (for example, keras.ops.matmul or keras.ops.softmax) instead of TensorFlow ops, JAX arrays, or PyTorch tensors directly. The active backend is selected with the KERAS_BACKEND environment variable or keras.config.set_backend().

Version assumptions: Keras 3 requires Python 3.9 or newer, plus TensorFlow 2.15+, JAX 0.4.25+, or PyTorch 2.1+ depending on the backend. Older environments may need dependency upgrades that break legacy tf.keras code.

Worked example: a backend-agnostic custom layer

This layer scales and shifts its input using only keras.ops, so the same class definition works on any supported backend.

import keras
from keras import layers, ops


class ScaleAndShift(layers.Layer):
    def __init__(self, scale=1.0, shift=0.0, **kwargs):
        super().__init__(**kwargs)
        self.scale = scale
        self.shift = shift

    def call(self, inputs):
        return ops.multiply(inputs, self.scale) + self.shift


model = keras.Sequential([
    layers.Input(shape=(784,)),
    ScaleAndShift(scale=2.0, shift=-1.0),
    layers.Dense(10, activation="softmax"),
])

model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.save("portable_model.keras")

To verify portability, initialize the model on one backend, then load it on another:

# Terminal 1: TensorFlow backend
export KERAS_BACKEND=tensorflow
python -c "import keras; m = keras.saving.load_model('portable_model.keras'); print(keras.backend.backend())"

# Terminal 2: JAX backend
export KERAS_BACKEND=jax
python -c "import keras; m = keras.saving.load_model('portable_model.keras'); print(keras.backend.backend())"

Where: a terminal with the Keras environment activated. Permissions: none required for reading the model file. Expected check: each command prints the backend name matching the environment variable. Risk: the .keras archive stores weights and metadata, but numerical results may differ slightly between backends because of floating-point promotion and kernel implementation. Compare inference outputs with a tolerance such as numpy.allclose(..., atol=1e-4) rather than expecting exact equality.

Where portability breaks down

• Performance is not guaranteed to match. JAX often leads on TPU and GPU workloads because of JIT compilation, while the PyTorch backend can lag on dynamic shapes. Benchmark your model and hardware before committing to a backend.
• Custom training loops using backend autodiff directly are not portable. Avoid tf.GradientTape or torch.autograd in shared code; use the Keras optimizer and training APIs instead.
• Some callbacks are backend-specific. TensorBoard profiling and PyTorch Lightning integrations may only work on their native backend. Verify callback compatibility before switching.
• Legacy tensorflow.keras imports and custom TensorFlow ops may not load on JAX or PyTorch. Audit imports and migrate to the unified keras namespace.
• Debugging can be opaque because stack traces originate in backend kernels. During development, disable Keras traceback filtering to see full traces.

Pre-flight checklist before switching backends

1. Confirm the active backend with python -c "import keras; print(keras.backend.backend())".
2. Replace backend-specific tensor operations in custom layers with keras.ops equivalents.
3. Save the model with model.save("model.keras") and load it on the target backend to confirm the archive reconstructs correctly.
4. Run inference on a fixed input on both backends and compare outputs within a numerical tolerance.
5. Profile step time and memory on the target hardware; do not assume parity with the source backend.

Bottom line

Keras 3 makes it possible to write one model and run it on TensorFlow, JAX, or PyTorch, but portability depends on staying inside the Keras API surface. Treat keras.backend.* as an escape hatch, not a default. Isolate any backend-specific code in small, tested utilities, and verify with a cross-backend load and inference check before you change production training or serving pipelines.