C.W.K.
Stream
Lesson 04 of 07 · published

Cross-Backend Saving

~9 min · serialize

Level 0Keras Apprentice
0 XP0/97 lessons0/20 achievements
0/120 XP to next level120 XP to go0% complete

Train under one engine, load under another

This is a Keras 3 superpower: save a model under one backend and load it under a completely different one. Train on JAX to get its fast JIT compilation, save the .keras file, then reopen it under TensorFlow to hand off to TF Serving — the file doesn't care which engine wrote it.

Why it works

The .keras format stores no backend-specific operations. It records the layer config and the weights as plain, backend-agnostic arrays — there's no traced JAX graph or frozen TF function baked in. So when a different backend loads the file, it rebuilds the same layers from the config and dispatches their math through its own keras.ops implementation. Numerically identical model, different engine underneath.

Where portability breaks

The guarantee holds as long as you stay inside Keras. If a layer reaches past keras.ops into a backend's native API — a raw tf.signal.fft call, a hand-written PyTorch op — that operation has no equivalent on the other backends, and the cross-backend load will fail there. The rule is simple: build models out of Keras layers and keras.ops, and portability is free; drop to native ops and you've pinned yourself to one engine.

Code

Save under JAX, reload under TensorFlow, export for serving·python
# Train on JAX (fast JIT compilation)
os.environ["KERAS_BACKEND"] = "jax"
model.fit(x_train, y_train, epochs=10)
model.save("trained_on_jax.keras")

# Later: load on TensorFlow for deployment
os.environ["KERAS_BACKEND"] = "tensorflow"
model = keras.models.load_model("trained_on_jax.keras")
model.export("tf_serving_model", format="tf_saved_model")

External links

Exercise

Train MNIST with KERAS_BACKEND=tensorflow. Save as mnist.keras. In a new shell with KERAS_BACKEND=torch, load and predict on test set. Verify accuracy matches the original training run.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.