Keras 3's Multi-Backend Architecture: Write Once, Run on TensorFlow, JAX, or PyTorch — With Caveats

The Framework Lock-In Problem

Deep learning teams often pick a framework early — TensorFlow for production tooling, PyTorch for research flexibility, JAX for TPU performance — and then find themselves stuck. Moving a model between frameworks typically means rewriting training loops, custom layers, and distribution logic. Keras 3 (released 2024) changes this equation by decoupling the high-level API from the execution engine, but the portability promise comes with boundaries that matter in practice.

How the Backend Abstraction Works

Keras 3 introduces a backend-agnostic layer that handles tensor operations, automatic differentiation, compilation, and device placement. You select the backend at runtime via the KERAS_BACKEND environment variable or keras.config.set_backend(). The same model definition — whether built with the Functional API or by subclassing Layer/Model — can run on TensorFlow, JAX, or PyTorch without code changes provided you stay within the Keras-native API surface.

Under the hood, each backend maps Keras operations to its native primitives: XLA compilation for JAX and TensorFlow, TorchCompile for PyTorch. Keras-native features like Model.fit, callbacks, and the .keras serialization format work across backends, though numerical behavior (RNG seeding, gradient accumulation order, bfloat16 precision on TPU vs GPU) can differ.

A Worked Example: Switching Backends in Practice

Create a minimal script to verify portability. Save this as train_mnist.py:

import os
import keras
from keras import layers, models, datasets

# Load data
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.0
x_test = x_test.reshape(-1, 28, 28, 1).astype('float32') / 255.0

# Model definition — pure Keras, no backend-specific ops
inputs = keras.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, activation='relu')(inputs)
x = layers.MaxPooling2D()(x)
x = layers.Conv2D(64, 3, activation='relu')(x)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(10, activation='softmax')(x)
model = models.Model(inputs, outputs)

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

print(f"Backend: {keras.config.backend()}")
model.fit(x_train, y_train, epochs=2, batch_size=128, verbose=1)
model.save('mnist_model.keras')
print("Model saved to mnist_model.keras")

Run it on each backend from your terminal (requires pip install keras>=3.0 plus the respective backend package):

# TensorFlow backend (default if tensorflow installed)
KERAS_BACKEND=tensorflow python train_mnist.py

# JAX backend
KERAS_BACKEND=jax python train_mnist.py

# PyTorch backend
KERAS_BACKEND=torch python train_mnist.py

Each run prints the active backend and trains the same model. The saved mnist_model.keras (ZIP-based JSON config + NumPy weights) can be loaded and run for inference on another backend if the target backend is installed and the model contains no backend-specific operations. Try loading it in a separate script with a different KERAS_BACKEND to confirm.

Trade-Offs and Limitations

Portability holds when you stay within the keras.ops namespace and Keras-native layers. Backend-specific operations — tf.nn.atrous_conv2d, jax.lax.conv_general_dilated, torch.nn.functional — break cross-backend compatibility. Custom layers using tf.TensorArray, jax.lax.scan, or PyTorch autograd internals will not transfer.

Distribution strategies differ: TensorFlow uses tf.distribute.Strategy subclasses, JAX uses jax.sharding.Mesh with PJit, PyTorch uses FSDP. Keras 3 abstracts Model.fit for basic data parallelism, but advanced sharding requires backend-specific code. Mixed-precision behavior also varies — bfloat16 on TPU (JAX) versus GPU (PyTorch/TensorFlow) can produce numerical differences.

Serialization favors the native .keras format. SavedModel export remains TensorFlow-only. ONNX conversion from JAX or PyTorch backends is less mature than TensorFlow's tf2onnx path (as of 2024; verify current operator support for your inference runtime).

Actionable Next Steps

1. Audit your codebase for hardcoded backend imports (import tensorflow as tf in model files). Replace with keras.ops equivalents or conditional imports guarded by keras.config.backend().
2. Run a parity benchmark: train a standard model (ResNet-50 on an ImageNet subset) on each backend with identical hyperparameters. Compare loss curves, step time, and final accuracy. Expect minor numerical variance.
3. Validate serialization round-trips: save on one backend, load and run inference on another. Check output numerical equivalence within a reasonable tolerance (e.g., np.allclose with rtol=1e-4).
4. Check third-party dependencies: older Keras Applications or community layers may hardcode TensorFlow. Verify Keras 3 compatibility before adopting.

Keras 3's multi-backend architecture makes framework portability a practical reality for many workflows — provided you respect its boundaries and verify the behavior that matters to your deployment target.