Why doesn’t a Flax model store its own weights?

Because a Flax Module is built to be a specification, not a container. Flax’s own developer notes put it plainly: constructing a Module "does not actually create variables or any internal state whatsoever. It’s best to think of it as a specification or template of the Module that contains functionality but no data." The weights live in a separate variables dict that init() hands back to you, and every later call — apply(), a checkpoint save, a jax.grad — passes that dict in and out explicitly. The model object is architecture; the dict is data. Keeping them apart is what "params are data" means.

What is the difference between init and apply in Flax?

Less than it looks. Flax’s module-lifecycle notes are direct about it: "there actually is no separate initialization path in Flax… init is nothing more than a wrapper around apply," calling your module with an empty variables dict, a PRNG stream named 'params' supplied automatically, and every collection marked mutable so any parameter you touch gets created on the spot. apply() is the general case: you hand it a real variables dict, it runs your code against those values, and by default nothing is mutable — it just computes. init() is apply() the first time, before there is anything to compute with.

What is a params pytree in Flax and why does it matter?

It is the plain nested dictionary that init() returns — collection name at the top ("params"), then module hierarchy underneath, so a two-layer MLP comes back as {"params": {"Dense_0": {"kernel": ..., "bias": ...}, "Dense_1": {"kernel": ..., "bias": ...}}}. Nothing about that structure is special-cased; it is a JAX pytree, the same kind of nested container jax.tree_util.tree_map already knows how to walk. That is the entire payoff of keeping params outside the model: an optimizer like Optax, a checkpoint library like Orbax, a sharding annotation — none of them need to understand "Flax Module." They just need a pytree, and a pytree is exactly what they get.

How does Flax thread PRNG keys through nested modules without passing them by hand?

By folding two things into a hash every time self.make_rng is called: the calling module’s path in the tree (self.scope.path — empty at the root, 'Dense_0' one level down, and so on) and how many times that specific module has drawn from that specific stream before. The resulting integer gets folded into the stream’s starting key with jax.random.fold_in, producing a fresh, unique key. Because every submodule keeps its own call count per stream, the same top-level key reproduces the identical split at every leaf on every run — you hand init() one key, and the tree does the rest deterministically.

What is the conditional-submodule bug in Flax’s @nn.compact style?

It happens when a submodule’s auto-generated name depends on which branch of an if/elif actually ran. Inside @nn.compact, an unnamed submodule is named by construction order — the first Dense becomes "Dense_0" regardless of which code path built it. Flax’s own module-lifecycle note gives the exact failure: a method that returns nn.Dense(8)(x) in an "encode" branch and nn.Dense(4)(x) in a "decode" branch "will break because either the encoder or decoder path will construct a Module named 'Dense_0'. This means the two Modules will share parameters which is not intended here." The fix is to give each branch an explicit name, build the submodule before the branch, or assign both submodules up front in setup(), where they exist as named attributes no matter which one __call__ actually uses.

Params Are Data: the Functional Model Behind Flax

A Flax model never holds its own weights. Call init() and it hands back an ordinary nested dict of arrays; call apply() and it takes that same dict back in to run the forward pass. We follow one PRNG key as it splits down the module tree, watch every self.param drop a leaf into the growing params pytree, send that pytree back through apply, and catch the exact bug Flax’s own docs name when submodules get built inside a branch — the discipline that lets checkpointing, sharding, and jax.grad treat a whole model like any other array.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Params Are Data: the Functional Model Behind Flax — built to be understood, not skimmed.

Questions

Why doesn’t a Flax model store its own weights?: Because a Flax Module is built to be a specification, not a container. Flax’s own developer notes put it plainly: constructing a Module "does not actually create variables or any internal state whatsoever. It’s best to think of it as a specification or template of the Module that contains functionality but no data." The weights live in a separate variables dict that init() hands back to you, and every later call — apply(), a checkpoint save, a jax.grad — passes that dict in and out explicitly. The model object is architecture; the dict is data. Keeping them apart is what "params are data" means.
What is the difference between init and apply in Flax?: Less than it looks. Flax’s module-lifecycle notes are direct about it: "there actually is no separate initialization path in Flax… init is nothing more than a wrapper around apply," calling your module with an empty variables dict, a PRNG stream named 'params' supplied automatically, and every collection marked mutable so any parameter you touch gets created on the spot. apply() is the general case: you hand it a real variables dict, it runs your code against those values, and by default nothing is mutable — it just computes. init() is apply() the first time, before there is anything to compute with.
What is a params pytree in Flax and why does it matter?: It is the plain nested dictionary that init() returns — collection name at the top ("params"), then module hierarchy underneath, so a two-layer MLP comes back as {"params": {"Dense_0": {"kernel": ..., "bias": ...}, "Dense_1": {"kernel": ..., "bias": ...}}}. Nothing about that structure is special-cased; it is a JAX pytree, the same kind of nested container jax.tree_util.tree_map already knows how to walk. That is the entire payoff of keeping params outside the model: an optimizer like Optax, a checkpoint library like Orbax, a sharding annotation — none of them need to understand "Flax Module." They just need a pytree, and a pytree is exactly what they get.
How does Flax thread PRNG keys through nested modules without passing them by hand?: By folding two things into a hash every time self.make_rng is called: the calling module’s path in the tree (self.scope.path — empty at the root, 'Dense_0' one level down, and so on) and how many times that specific module has drawn from that specific stream before. The resulting integer gets folded into the stream’s starting key with jax.random.fold_in, producing a fresh, unique key. Because every submodule keeps its own call count per stream, the same top-level key reproduces the identical split at every leaf on every run — you hand init() one key, and the tree does the rest deterministically.
What is the conditional-submodule bug in Flax’s @nn.compact style?: It happens when a submodule’s auto-generated name depends on which branch of an if/elif actually ran. Inside @nn.compact, an unnamed submodule is named by construction order — the first Dense becomes "Dense_0" regardless of which code path built it. Flax’s own module-lifecycle note gives the exact failure: a method that returns nn.Dense(8)(x) in an "encode" branch and nn.Dense(4)(x) in a "decode" branch "will break because either the encoder or decoder path will construct a Module named 'Dense_0'. This means the two Modules will share parameters which is not intended here." The fix is to give each branch an explicit name, build the submodule before the branch, or assign both submodules up front in setup(), where they exist as named attributes no matter which one __call__ actually uses.

Params Are Data: the Functional Model Behind Flax

Questions

Related explainers