Diagnose partial initial model states and document Axon.build

[seanmor5]

Closes #475.

Most of the concrete asks in that thread have already landed: the Axon.CompileError message now points at debug: true for the "layer was defined at" section, Axon.block and Axon.namespace/2,3 give parameter maps grouping/nesting, and Axon.elem/3 / Axon.fetch/3 destructure containers. The Livebook table truncation and the unpickler location were answered in the thread as not being Axon changes. Two asks were still open, and this PR resolves them.

Initial-state diagnostics in init_fn

The reporter asked for stricter modes that flag parameters that were missing from the supplied parameter map (and therefore silently initialized) and extra unused data in the map. The maintainer position in the thread was not to raise, because partially initializing a model is a normal training / transfer-learning workflow, but to do what Bumblebee's loader does and log the missing parameters at debug level, and to log unused parameters as well.

init_fn now does both, and always names parameters by their full dotted path:

• When the initial state is not empty, every parameter that was absent from it and had to be initialized from the layer's initializer is listed in one Logger.debug line. With Axon.ModelState.empty() nothing is logged, since every parameter is "missing" and the line would be noise.
• Keys in the initial state that do not correspond to any model parameter are collected into one Logger.warning line with their full paths. Previously this was one warning per key with only the leaf name, so a naming mismatch in a large model produced hundreds of found unexpected key ...: "kernel" lines that did not say which layer was affected. An entirely unknown subtree is reported by its root ("dense_9") rather than by every leaf under it.
• The existing "initial type for parameter ... does not match policy" warning now also names the full path.

Paths follow the structure of Axon.ModelState.data, so Axon.block sub-states appear as "block_0.dense_0.kernel" and composite parameters as "lstm.input_kernel.wii". Quantized tensors and shared parameters are structs and are treated as leaves.

For example, loading a state that misspells one key and covers only the first layer:

model =
  Axon.input("input", shape: {nil, 1})
  |> Axon.dense(2, name: "dense_0")
  |> Axon.dense(2, name: "dense_1")

{init_fn, _} = Axon.build(model)
init_fn.(Nx.template({1, 1}, :f32), Axon.ModelState.new(%{"dense_0" => %{"kernel" => k, "bais" => b}}))

logs

[debug] the following parameters were not present in the initial model state and were initialized from scratch: "dense_0.bias", "dense_1.bias", "dense_1.kernel"
[warning] the following keys in the initial model state do not correspond to any parameter of the model and were ignored: "dense_0.bais"

Tradeoffs: the missing-parameter report is at :debug rather than :warning or an exception, following the maintainer's comment; users validating a ported checkpoint can raise their log level. The check lives in init_fn (the natural point where the supplied state meets the model's declared parameters) rather than in predict_fn, and the Axon.build docs now recommend running init_fn with loaded parameters as the way to validate them. No new Axon.build option was added. Because init_fn is JIT-compiled and cached by argument signature, a repeated call with an identical signature may not re-trace and therefore may not re-log; this was already true of the existing warnings. Merge semantics are unchanged: this only changes what is logged.

Axon.build/2 docs and guides

The reporter also asked whether predict_fn mutates backend state or whether init_fn copies parameters to the backend, and proposed wording for the docs. The init_fn / predict_fn sections of Axon.build/2 now explain that both are pure functions of their arguments, that the model never carries parameters and nothing is cached on the backend, why a template is needed, what init_fn does with an initial state (use, cast, initialize, ignore-with-warning), and what predict_fn returns in :inference versus :train mode. The old example passed a bare map to init_fn, which now triggers a deprecation warning; it uses Axon.ModelState.new/1 and Axon.ModelState.empty/0 instead. The 22 init_fn.(x, %{}) call sites in the guides, the three in the Axon moduledoc and the one in the README were updated the same way, and the signature description in the "Your first Axon model" guide now uses Axon.ModelState.t(). Livebook output blocks were left untouched.

Out of scope

A scale-factor form of Axon.resize for inputs with nil spatial dimensions was raised late in the thread and never answered by the maintainers; it is a separate layer feature and not addressed here. A predict-time unused-key check is also not added, since init_fn is where the supplied state is reconciled with the model.

Tests

test/axon/compiler_test.exs gains an "initial state diagnostics" describe block covering: the debug line for a partial state (and that present parameters are not listed), no debug line for an empty state, a single aggregated warning with full paths for unexpected keys (including an unknown subtree reported by its root), nested paths through Axon.block, nested paths through a composite LSTM parameter, and the full path in the type-cast warning. Five of the six fail on main. mix test: 883 passed (200 doctests, 683 tests), 47 excluded.

🤖 Generated with Claude Code