Choosing a Keras 3 Backend: How to Benchmark JAX, TensorFlow, and PyTorch on Your Own Model

Keras 3 runs the same Model on JAX, TensorFlow, or PyTorch by setting one environment variable. That is a genuinely useful property, and it invites a bad question: "which backend is fastest?" There is no portable answer. Throughput depends on your ops, your input pipeline, your batch size, your precision policy, and your hardware. What there is, is a cheap measurement — usually half a day — that answers it for your model.

This post is the harness we run on modernization engagements after a Keras 3 migration, and the non-performance constraints that usually decide the choice before the numbers arrive.

Switch the backend correctly first

The backend is selected before Keras is imported:

import os
os.environ["KERAS_BACKEND"] = "jax"   # or "tensorflow", "torch"
import keras

Setting it after import keras does nothing, and a stale ~/.keras/keras.json silently overrides what you thought you configured. Print keras.backend.backend() at the start of every run and log it in the run record. Half the "JAX is slower for us" reports we chase turn out to be TensorFlow runs.

If your model uses keras.ops and Keras layers throughout, the switch is free. If it calls tf.* inside a custom layer, that code is TensorFlow-only and must be rewritten in keras.ops before any comparison is meaningful. Backend portability is a property you earn in the layer code, not a flag.

Separate the input pipeline from the model

This is the step most comparisons skip, and it invalidates the rest. tf.data works as an input pipeline under all three backends — Keras converts batches at the boundary — so you can hold data loading constant while you vary the compute backend. Do that. Comparing a tf.data pipeline under TensorFlow against a PyTorch DataLoader under Torch measures your two loaders, not the backends.

Before timing anything, confirm you are not input-bound. Run a few hundred steps over a cached, pre-batched in-memory tensor. If step time barely changes versus the real pipeline, you are compute-bound and the comparison is valid. If it drops sharply, you are feeding-bound: fix the pipeline first (see our tf.data checklist), because until you do, all three backends will benchmark identically and you will conclude "no difference" for the wrong reason.

A fair harness

Hold everything constant except KERAS_BACKEND:

• Same model code, same initialization seed, same batch size, same precision policy (keras.mixed_precision.set_global_policy("mixed_float16") applies across backends — set it identically or not at all).
• Same machine, same driver, one run at a time. No shared GPU.
• Discard warmup. The first steps include compilation: XLA tracing under JAX, tf.function retracing under TensorFlow, torch.compile if you enabled it. Throw away at least 20–50 steps.
• Then time 200–500 steps and report median step time, plus p90. A mean over a run with one 40-second compile stall tells you nothing.
• Record peak device memory alongside step time. A backend that is 8% faster but leaves no headroom to raise batch size is often the worse choice.

import time
times = []
for i, batch in enumerate(ds):
    t0 = time.perf_counter()
    metrics = model.train_on_batch(*batch)
    if i >= warmup:
        times.append(time.perf_counter() - t0)

Under JAX, dispatch is asynchronous; a naive timer around an async call measures queueing, not compute. Force completion each step (block on the returned values) or time a whole epoch and divide. The same caution applies to CUDA streams under Torch — synchronize before you stop the clock.

Measure convergence, not just speed

Step time is the easy half. Run each backend to the same number of epochs on the same split and compare the loss curves and final validation metric across three seeds. You are looking for one thing: agreement. Small differences are expected — kernel accumulation order, default epsilon handling in some ops, and fused-kernel differences all move the fourth decimal place. A visible divergence in the curve is a bug, not a backend property, and it is almost always a custom layer, a custom loss, or a numerically fragile op reading differently on the new backend. Find it before you decide anything.

A backend that is 20% faster per step but needs more epochs to reach the target metric is not faster. Report time to target metric, which is the number the business actually pays for.

What usually decides it anyway

In practice we have seen the choice settle on constraints more often than on step time:

• Serving path. If the artifact ships through TF Serving, LiteRT, or a TFLite-based edge stack, training on the TensorFlow backend removes an export conversion from the pipeline. That is worth more than a few percent throughput.
• Ecosystem debt. Existing TFRecord stores, TFX pipelines, or tf.distribute strategies favor TensorFlow. Existing PyTorch tooling — vendor inference SDKs, in-house Torch datasets, or teammates who read Torch fluently — favors Torch.
• Scale-out shape. JAX's sharding and XLA compilation are strong when you are scaling across many accelerators or onto TPUs, and JAX tends to show its advantage on models that compile cleanly into large fused graphs. It gives up ground when your model has heavy dynamic shapes or Python-side control flow that forces recompilation.
• Compile time. XLA compilation cost is paid once per shape signature. On a long training run it is noise; in a short-iteration research loop, or with variable sequence lengths recompiling constantly, it is not. Pad or bucket shapes if you go this way.
• Debuggability. Eager Torch is the easiest to step through with a debugger. That matters during development, and matters less once training is stable.

What to do with the result

Write it down. A one-page table — backend, median step time, peak memory, time to target metric, and the constraint notes — ends the recurring hallway argument and gives the next engineer the reasoning, not just the conclusion. Re-run it when the model architecture changes materially, when you change hardware, or after a major Keras or backend release. It is a cheap measurement; the value is that it replaces a preference with a number.

And keep the model backend-portable regardless of which one wins. The point of writing layers in keras.ops is not that you switch often — it is that switching stays a half-day experiment instead of a rewrite when hardware, pricing, or a serving requirement changes underneath you.