One Model, Three Backends: Keras 3 on TensorFlow, JAX, and PyTorch

Keras used to be "the high-level API for TensorFlow". Since Keras 3 it is a standalone library that runs the same model code on TensorFlow, JAX, or PyTorch, with an OpenVINO backend for inference only. For a consulting firm that has spent years in TensorFlow, this is the single most useful change in the ecosystem: we can write one model and put it on whichever backend fits the client's training infrastructure and serving stack. This tutorial shows how it works in practice and when it is worth the trouble.

Choosing a backend

The backend is selected once, at import time, with the KERAS_BACKEND environment variable:

KERAS_BACKEND=jax python train.py
KERAS_BACKEND=tensorflow python train.py
KERAS_BACKEND=torch python train.py

or, in code, before import keras:

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras
print(keras.backend.backend())   # "jax"

A ~/.keras/keras.json file can also set it. Whatever you choose, set it explicitly in CI and in your training containers; "whichever backend happened to be installed" is not a configuration.

Install the backend you need alongside Keras - for example pip install keras jax[cuda12] or pip install keras tensorflow. Keras 3 is also what tf.keras resolves to in TensorFlow 2.16 and later, so a plain TensorFlow install already has it.

Writing backend-agnostic code

The rule is simple: use keras.ops, not tf.*, jnp.*, or torch.*, inside your model code. keras.ops implements the NumPy API plus the neural-network ops you need, and dispatches to the active backend.

A custom layer written this way runs everywhere:

import keras
from keras import ops, layers


class RMSNorm(layers.Layer):
    def __init__(self, epsilon=1e-6, **kwargs):
        super().__init__(**kwargs)
        self.epsilon = epsilon

    def build(self, input_shape):
        self.scale = self.add_weight(
            shape=(input_shape[-1],), initializer="ones", name="scale"
        )

    def call(self, x):
        variance = ops.mean(ops.square(x), axis=-1, keepdims=True)
        return x * ops.rsqrt(variance + self.epsilon) * self.scale


inputs = keras.Input(shape=(128,))
x = layers.Dense(256, activation="gelu")(inputs)
x = RMSNorm()(x)
outputs = layers.Dense(10, activation="softmax")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adamw", loss="sparse_categorical_crossentropy", metrics=["accuracy"])

Nothing in that file mentions a backend. The same is true of losses, metrics, optimizers, and callbacks from keras.*.

The places where backend-specific code sneaks in:

• Data pipelines. tf.data works as input to model.fit on every backend (Keras converts batches), so you can keep your existing pipeline. PyTorch DataLoader objects also work on every backend. What does not work is mixing tensors from the wrong framework inside the model.
• Custom training loops. model.fit is portable. A hand-written loop is inherently backend-specific - if you need one, write it with the backend's own primitives and keep it in a separate file.
• Random numbers. Use keras.random with a keras.random.SeedGenerator, not tf.random.
• Saving. Save in the .keras format. It is backend-neutral and loads on any backend.

Train on JAX, serve on TensorFlow

This is the combination we use most. JAX is usually the fastest backend for training on GPUs and TPUs, and TensorFlow has the most mature serving story (TF Serving, SavedModel, LiteRT). Keras 3 lets you use both:

# train.py  (KERAS_BACKEND=jax)
model.fit(train_ds, validation_data=val_ds, epochs=10)
model.save("model.keras")

# export.py  (KERAS_BACKEND=tensorflow)
import keras
model = keras.saving.load_model("model.keras")
model.export("saved_model/1")        # TensorFlow SavedModel for TF Serving

The .keras file carries weights and architecture; the TensorFlow process rebuilds the model on its own backend and exports a SavedModel with a serving_default signature. From there it is a standard TF Serving or Vertex AI deployment - see the TF Serving guide - and you can feed the same SavedModel to the LiteRT converter for on-device use.

Before you rely on this, verify numerics once: run a fixed batch through the model on both backends and assert_allclose the outputs. Backends differ in reduction order and default precision, so expect small float differences; anything larger than that means a custom layer is doing something non-portable.

Distributed training

For multi-GPU or TPU training on the JAX backend, Keras 3 provides keras.distribution with data-parallel and model-parallel layouts:

import keras
devices = keras.distribution.list_devices()
keras.distribution.set_distribution(keras.distribution.DataParallel(devices=devices))
# build + fit as usual

The Keras distributed training guide covers the model-parallel case, which is what you need for fine-tuning larger open-weight models.

OpenVINO for inference

The OpenVINO backend is inference-only: you cannot train on it, but you can load a .keras model and call predict with Intel's optimized CPU runtime underneath. It is worth knowing about for CPU-only serving environments where you do not want to carry a full training framework.

Is multi-backend worth it for your team?

Yes, when:

• Training and serving have different owners or different infrastructure (the JAX-train / TF-serve case above).
• You maintain model code that must run in both a PyTorch-native research group and a TensorFlow-native production group.
• You want to benchmark backends for a specific workload before committing - it is an environment variable, not a rewrite.

No, or not yet, when:

• Your codebase is full of tf.* calls inside layers and you have no budget to convert them to keras.ops. You can still run Keras 3 on the TensorFlow backend with tf.* inside layers; you just lose portability, and that is fine.
• You rely on custom training loops with framework-specific tricks.
• Your dependency on Keras 2 behaviours is deep; tf-keras exists as a bridge, but plan the move rather than living there.

The practical advice we give clients: write new layers with keras.ops, save in .keras format, and keep tf.data for input. That costs nothing today and keeps the backend decision open.

Need help moving a TensorFlow codebase onto Keras 3, or proving parity across backends? Contact us - it is a large part of our modernization work.