The line between portable and leaky
Knowing exactly what survives a backend switch and what doesn't is the difference between trusting the abstraction and being surprised by it. The good news: almost everything you write as Keras code is invariant. The differences only surface at the seam where you touch backend-native objects directly.
What stays the same
- Model definition — Sequential, Functional, and Subclassing all identical
- The training loop —
model.compile()+model.fit() - Callbacks — EarlyStopping, ModelCheckpoint, and the rest
- The
.kerassaving format - Every
keras.opsoperation
What changes (and where it leaks)
The differences below only bite when you reach past Keras to the engine — pulling a raw tensor out, placing it on a device by hand, or mixing in a backend-native library:
| Aspect | TensorFlow | PyTorch | JAX |
|---|---|---|---|
| Tensor type | tf.Tensor | torch.Tensor | jax.Array |
| Device mgmt | tf.device | .to(device) | jax.devices() |
| Layers are | Keras layers | torch.nn.Module | Keras layers |
| JIT compilation | tf.function | torch.compile | jax.jit (default) |
The practical rule: stay inside Keras + keras.ops and none of this column-shopping matters; the model is genuinely portable. The moment you cross into native code, mark that boundary clearly — it's the one place a backend swap can break you.