LayerGenius | Migrating tf.keras Code to Keras 3: What Actually Breaks

Keras 3 is a genuinely different library from the tf.keras that shipped inside TensorFlow for years: one API, three backends (TensorFlow, JAX, PyTorch), and a set of deliberate breaks with old habits. Most model definitions port with minor edits. The pain concentrates in a handful of predictable places. This post is the checklist we work through when moving a legacy tf.keras codebase.

Start with the import line

Keras 3 is a standalone package. The migration begins by replacing

from tensorflow import keras
from tensorflow.keras import layers

with

import keras
from keras import layers

If your codebase mixes both styles — and most older ones do — normalize first, in a commit of its own, while still on Keras 2. tf.keras and a separately installed keras package must never both be imported in one process; that configuration produces two different class hierarchies and baffling isinstance failures. If you need to stay on the old behavior temporarily, pin the tf_keras compatibility package and set TF_USE_LEGACY_KERAS=1, then migrate deliberately rather than by accident.

The big one: raw TensorFlow ops in model code

The most common source of breakage is TensorFlow ops embedded in custom layers, losses, and metrics. Code like tf.reduce_mean, tf.einsum, or tf.where inside a call() method ties the model to the TF backend and, in several cases, fails outright under Keras 3's tracing.

The fix is keras.ops, which mirrors most of the NumPy API plus the common neural-network ops:

def call(self, x):
    m = keras.ops.mean(x, axis=-1, keepdims=True)
    return keras.ops.where(x > m, x, keras.ops.zeros_like(x))

Rewriting op calls is mechanical; finding them is the work. Grep for tf. inside every subclassed layer, model, loss, metric, and callback, and treat each hit as a decision: port to keras.ops (keeps the code backend-portable) or declare the component TF-only on purpose.

Lambda layers deserve special suspicion. A Lambda wrapping an arbitrary TF closure often survives tracing but breaks serialization. Almost every Lambda worth keeping is clearer as a small subclassed layer with a registered serialization config.

`keras.backend` is gone

Legacy codebases lean on K.function, K.learning_phase, K.int_shape, and friends. That namespace no longer exists in any useful form. Most uses map to keras.ops (shapes, math) or disappear entirely (learning phase is handled by the training argument). Budget real time here if your code predates TF 2 — K.* usage correlates strongly with age.

Saving and loading

Keras 3's native format is the .keras zip archive, and it is strict: custom objects must be registered (@keras.saving.register_keras_serializable()) and get_config() must actually round-trip. Legacy H5 files still load for most architectures, but SavedModel-format models do not load back as Keras models. For serving, that path changed name rather than disappearing: model.export("path") produces a SavedModel with an inference-only signature for TF Serving. Keep training checkpoints in .keras, and treat exported artifacts as deployment outputs, not something you reload for further training.

Custom training loops

If you overrode train_step, the port is usually straightforward under the TF backend — the GradientTape pattern still works. But if the reason for adopting Keras 3 is the JAX backend, train_step must be rewritten in the stateless style: JAX requires pure functions, so state (trainable variables, optimizer slots, metrics) is threaded through explicitly via compute_loss_and_updates and stateless_apply. This is the deepest change in the migration; plan it as its own task, not a find-and-replace.

Two smaller traps in the same area:

RNG. tf.random calls inside model code should become keras.random with an explicit keras.random.SeedGenerator — required for JAX, healthier everywhere.
Compiled step defaults. Keras 3 is more aggressive about compilation (XLA under JAX, jit_compile under TF). Ops with data-dependent shapes that ran fine eagerly can fail under compilation; the error messages point at the tracer, not your line. Bisect by temporarily passing jit_compile=False to compile().

The input pipeline can stay

Good news: tf.data pipelines work as input to Keras 3 on all backends — a JAX- or Torch-backed model will happily consume a tf.data.Dataset. You do not need to rewrite data loading to migrate. Preprocessing layers inside the model (Normalization, StringLookup, and friends) are backend-portable; preprocessing done with raw tf.* ops in a Dataset.map remains fine because it runs in the pipeline, not the model.

A migration order that works

Normalize imports and pin versions; get the test suite green on Keras 2 first.
Inventory custom code: subclassed layers/models/losses/metrics/callbacks, Lambda layers, K.* usage, custom train_steps.
Port pure model-definition code to keras.ops; keep the TF backend while doing it so you change one variable at a time.
Re-establish save/load round-trips in .keras format, with serialization tests for every custom object.
Only then flip KERAS_BACKEND and chase backend-specific failures.
Run a short training job on a fixed seed and dataset slice, and compare loss curves against the pre-migration baseline before trusting anything at scale.

The codebases that migrate painfully are the ones that skip step 2 and discover their custom-code inventory one stack trace at a time. The inventory turns an open-ended debugging session into a bounded engineering task — which is what a migration should be.

Migrating tf.keras Code to Keras 3: What Actually Breaks