Keras spent a decade as TensorFlow’s friendly face. Keras 3 tore that coupling out: the same keras.Sequential, the same model.fit(), the same .keras file now run on TensorFlow, JAX, or PyTorch, selected with a single environment variable. That promise is only believable if the abstraction lives somewhere real, and it does — in keras/src, the package this tour walks through. If you have followed the earlier source dives into PyTorch and JAX, you have seen single-framework cores with their own dispatch and their own file formats. Keras is the thin, opinionated layer that sits above engines like those and lets one modeling script target any of them.

The tree is flat and readable. keras/src/backend holds the switch plus one subpackage per engine, keras/src/layers and keras/src/models define the object model, keras/src/trainers runs the loops, keras/src/ops is a NumPy-shaped math namespace, and keras/src/saving owns the file format. Every claim below links to the exact line on master, so you can read along without cloning anything.

The overview map below shows the whole story in one picture: config on the left picks a backend, layers and models stack on top of a common Operation base, the trainer consumes batches, and saving plus export carry the result out the door.

Overview of keras/src: runtime config feeds the backend switch, the Layer-Operation core supports Model, Functional and Sequential, and the Trainer plus saving modules form the outer ring.

Reading the overview from left to right:

• KERAS_BACKEND env override, L362 — after keras.json is parsed, this block in backend/config.py lets an environment variable override the stored backend name.
• backend switch, L34 — the star-import that pulls exactly one engine package into the keras.src.backend namespace.
• class Layer, L81 — Layer(BackendLayer, Operation): every layer is also an Operation, which is how symbolic tracing works for free.
• add_weight, L526 — the single funnel through which trainable state is allocated, dtype-policy aware.
• symbolic_call, L90 — in ops/operation.py, calls on KerasTensor inputs are recorded as graph nodes instead of executed.
• class Model, L38 — Model(Trainer, Layer), the multiple-inheritance hinge of the whole library.
• fit, L549 — the main loop in trainers/trainer.py, shared by Sequential, Functional and subclassed models.
• EpochIterator.next, L140 — normalizes NumPy arrays, tf.data, torch loaders and generators into per-step batches.
• save_model, L69 — writer for the .keras v3 zip format, with zip-bomb guards nearby at L429.

The detail map adds the four backend packages, the common runtime, the ops namespace, dtype policies, distribution, and the saving stack.

One Codebase, Three Backends

Selection is deliberately boring. Module-level defaults in backend/config.py set _BACKEND = "tensorflow", the BUILT_IN_BACKENDS frozenset at L18 names jax, numpy, tensorflow and torch, and PLUGGABLE_BACKENDS at L20 admits mlx, openvino and paddle through a plugin protocol. The public accessor backend() at L385 just returns the resolved string.

The interesting part is what happens next. backend/__init__.py runs a four-way if backend() == "tensorflow": ... chain and star-imports the winning package, so the names Variable, convert_to_tensor, compute_output_spec and friends are rebound to engine-specific implementations before anything else in Keras runs. Each engine package follows the same template — see the JAX variant and the torch variant, both of which re-export their own Variable and compute_output_spec. This is import-time polymorphism: no registry, no factory, just one branch taken once per process.

Above the switch sits a genuinely engine-agnostic contract. KerasTensor is a symbolic placeholder carrying shape, dtype and mask — no data. Variable in variables.py defines the stateful side with assign, assign_add and assign_sub at L282, each delegating to the backend value underneath. Mixed precision rides on DTypePolicy, which decides compute and variable dtypes per layer; quantized variants like QuantizedFloat8DTypePolicy at L388 extend the same class. Compared with the hand-written kernels you find in Triton or the fused CUDA work in FlashAttention, Keras deliberately owns none of that — it orchestrates.

Layers Are Operations

The inheritance line at layer.py L81 is the most important sentence in the codebase: a Layer is an Operation. That means the machinery in ops/operation.py — node tracking, compute_output_spec, symbolic dispatch — is inherited, not duplicated. When __call__ at L905 receives KerasTensor inputs, the layer does not run; it records a node and returns new symbolic tensors. When it receives real data, _dispatch_call executes call() normally.

The lifecycle is two-phase. build() at L430 runs once, the first time shapes are known, and every piece of state must be created through add_weight so dtype policies, distribution layouts and serialization metadata stay consistent. call() at L1104 is pure math against those weights.

Dense is the canonical example: class Dense(Layer) at L14, with build at L117 allocating kernel and bias lazily and call at L188 doing the matmul plus activation. At the other end of the complexity spectrum, MultiHeadAttention builds query, key and value projections in its build at L265 and consults the global flash-attention flag from config.py before choosing a kernel path in call at L611. The kernels it eventually hits live in ops/nn.py, where softmax and its Operation class sit side by side — every op in the namespace follows that class-plus-function pattern.

Functional Builds the Graph So You Do Not

keras.Sequential is the convenience wrapper — class Sequential(Model) at L20, with add at L78 growing the stack and building deferred until shapes are known. But the real modeling engine is Functional: class Functional(Function, Model). You hand it input and output tensors, and it walks the KerasTensor graph backwards, collecting every intermediate layer and the nodes that connect them. Because layers already record nodes via symbolic_call, Functional never re-runs anything to discover topology — the graph was written down while you were writing the model.

That same node graph is what makes .keras files small and reloads faithful: functional_from_config at L566 rebuilds the DAG from serialized layer configs and reconnects nodes in dependency order. Serialization round-trips through serialize_keras_object at L166 and its counterpart deserialize_keras_object at L405, which resolve classes from a registration table so custom layers survive a save-load cycle. For a serving-focused view of where these artifacts go, the tours of ONNX and Safetensors cover the formats Keras exports to or competes with.

From Install to Your First Model

Nothing exotic is required; Keras is plain PyPI:

pip install keras

By default you get the TensorFlow backend. To train on JAX or PyTorch instead, set the variable that config.py honors at L362 before the first import, for example KERAS_BACKEND=torch — the PyTorch tour covers what happens under that engine. Then a complete program:

import keras
from keras import layers

model = keras.Sequential([
    layers.Input(shape=(784,)),
    layers.Dense(64, activation="relu"),
    layers.Dense(10, activation="softmax"),
])
model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)
model.summary()
model.fit(x_train, y_train, epochs=3, batch_size=32)

Every symbol above resolves to source we have already visited: Sequential in models/sequential.py, Dense in layers/core/dense.py, compile at trainer.py L41, summary at model.py L222, and fit at L549. Inside the loop, EpochIterator accepts raw NumPy arrays, tf.data sets, torch DataLoaders or Python generators and yields uniformly shaped batches; CompileMetrics at L130 and CompileLoss at L429 aggregate the numbers even when your model has nested or dict-shaped outputs. Overriding train_step on a subclass replaces exactly one method and keeps the rest of the loop — that is the payoff of keeping the trainer a mixin on Model rather than a separate class hierarchy.

Saving, Loading and Exporting

The v3 format is a plain zip: config.json describing the architecture, metadata.json, a JSON weights manifest next to the weight data, and an assets/ folder for anything custom (the constants sit at saving_lib.py L44-L48). Model.save at L273 delegates to the public saving_api.save_model at L21, which dispatches to the v3 writer in saving_lib or to legacy h5 code for old files. The loader side, load_model at L127, rehydrates through the serialization registry and, in saving_lib, _save_state and _load_state at L744/L807 walk the object tree recursively so nested models, optimizers with slot variables, and metrics all round-trip.

For deployment, Model.export at L809 produces inference artifacts rather than training checkpoints; the ONNX path in export/onnx.py traces the model through whichever backend is live and writes a portable graph you can serve with the tools in the ONNX Runtime tour. There is also a distribution story: DeviceMesh at L155 and DataParallel at L438 shard variables and inputs across accelerates, and layer.py imports the layout machinery directly so add_weight can place weights on the mesh. For dictionary-shaped weights shared across frameworks, the approach rhymes with what Accelerate and DeepSpeed do at the trainer level.

Honest Limits

• The headline says three backends, and that undersells and oversells at once. There are four built-ins once you count the NumPy backend (useful for tests, too slow for real training), plus pluggable mlx, openvino and paddle targets that go through a protocol with less battle time than the big three.
• Backend parity is a spectrum, not a guarantee. The JAX and torch packages implement the same op surface, but numerical edge cases, dtype promotions and shape inference differ enough that compute_output_spec tracing occasionally disagrees with eager execution; tests in the repo explicitly compare all backends rather than asserting one truth.
• The v3 format is Keras-specific. It is open and documented, but unlike the ONNX story there is no cross-framework consumer; other tooling reads config.json only through Keras itself or reimplementations.
• fit is opinionated. Fancy data loading — streaming tokenization, on-the-fly augmentation on GPU, anything asynchronous — eventually fights EpochIterator, and at that point you subclass and override steps, which is supported but puts you back to writing framework code.
• Symbolic tracing has limits: dynamic sequence lengths, data-dependent control flow and Python conditionals inside call are fine eagerly but must be expressed carefully for jit_compile paths, especially on JAX.

None of these are disqualifying; they are the cost of the bargain Keras makes. It trades deep engine control for a stable, inspectable modeling surface — and after reading keras/src, the surface is no longer magic. Between the config file, the star-import switch, the Operation base and the zip on disk, roughly every behavior you trigger from user code has a one-file explanation. That is the best compliment available for an abstraction layer, and it is why Keras remains the shortest path from an idea to a trained, saved, exportable model — whether you saw this tour as a JAX user, a PyTorch loyalist, or someone shipping to OpenVINO at the edge.

Enjoyed this post? Never miss out on future posts by following us