Keras 3 — Multi-Backend Deep Learning • blog.dopana

Keras 3 is the multi-backend deep learning framework from keras-team — “deep learning for humans.” The key change from Keras 2: it no longer depends on TensorFlow. One codebase now runs on JAX, TensorFlow, PyTorch, and OpenVINO (inference-only). ~64k GitHub stars, used by 5M+ developers.

Why Multi-Backend Matters#

Historically you had to pick a framework and live with it. Keras 3 decouples the API from the compute substrate:

No lock-in — switch backends by changing one environment variable
Best performance — pick the fastest backend per model. Keras benchmarks show 20-350% speedups, with JAX typically winning on GPU/TPU/CPU
Ecosystem optionality — a Keras model is a PyTorch Module, can be exported as a TensorFlow SavedModel, or instantiated as a stateless JAX function

Choosing a Backend#

export KERAS_BACKEND="jax" # or: tensorflow, torch, openvino

bash

Or edit ~/.keras/keras.json. Backend minimum versions (Keras 3 stable): TensorFlow 2.16.1, JAX 0.4.20, PyTorch 2.1.0, OpenVINO 2025.3.

Backend	Best for
JAX	Fastest training/inference on GPU & TPU, XLA
TensorFlow	Existing TF stacks, tf.data pipelines
PyTorch	HF ecosystem, dynamic models, research
OpenVINO	CPU inference optimization (Intel)

keras.ops — One NumPy-Like API#

Keras 3 ships keras.ops, a full NumPy API that works on every backend — so custom layers, losses, and metrics run anywhere without tf. / torch. imports:

import keras
import numpy as np

x = keras.ops.ones((3, 3))
y = keras.ops.matmul(x, keras.ops.transpose(x))

python

Cross-Framework Data Pipelines#

Keras 3 models train with any pipeline — tf.data.Dataset, torch.utils.data.DataLoader, NumPy arrays, Pandas DataFrames, or keras.utils.PyDataset. No rewriting.

Full High-Level API#

All familiar pieces are here — layers, metrics, losses, optimizers, callbacks, training loops, saving/serialization — and they’re backend-agnostic. Custom train_step() must be written per backend, but compute_loss() works everywhere.

Example: Same Model, Any Backend#

import os
os.environ["KERAS_BACKEND"] = "jax"   # swap: "tensorflow", "torch"

import keras

model = keras.Sequential([
    keras.Input(shape=(28, 28)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dense(10, activation="softmax"),
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.fit(x_train, y_train, epochs=5)

python

Same code runs on JAX, TF, or PyTorch.

Pretrained Models#

Keras Applications: all 40 models (ResNet, EfficientNet, MobileNet…) on every backend
KerasCV & KerasHub: BERT, OPT, Whisper, T5, StableDiffusion, YOLOv8, SegmentAnything — all backends

Migration from Keras 2 / tf.keras#

Built-in-layer models migrate with near-zero changes. Custom tf.* code needs keras.ops replacements, and backend-specific train_step() overrides get per-backend implementations. After migration you can flip to JAX or PyTorch with one env var.

Pros & Cons#

Pros:

Write once, run on 4 frameworks
No framework lock-in — rare in 2026
SOTA performance via JAX backend
Backed by the same team (François Chollet) that built tf.keras

Cons:

Custom training loops need per-backend code
OpenVINO is inference-only
New distribution API (device mesh) is JAX-only for now
Niche custom ops may need backend-specific fallbacks

Conclusion#

Keras 3 is the closest thing the field has to a unified deep-learning API in 2026. It doesn’t pick a winner between TensorFlow, PyTorch, and JAX — it lets you write models once and run them anywhere, taking the best performance each framework offers. For teams building libraries or shipping to multiple users, it’s increasingly the default recommendation.

Keras 3 — Multi-Backend Deep Learning