Merge training updates into a single Axon.ModelState.update call

[seanmor5]

Closes #576.

Axon.Loop.train_step/4 updated the model state in two different ways. The objective function that runs under Nx.Defn.value_and_grad took the differentiated variable (the trainable parameters) and spliced it back into model_state.data by hand, using a private tree_merge/3 that was a copy of the one in Axon.ModelState. The step function then applied the optimizer output with the real API, Axon.ModelState.update/3. The hand-rolled merge existed because Axon.ModelState.update/3 used to be a plain def and could not be called from the jitted step. Since #656 turned the Axon.ModelState API into transforms, that restriction is gone.

This PR makes the objective call Axon.ModelState.update(model_state, trainable_parameters) directly and deletes the duplicated tree_merge/3 from Axon.Loop, so both parameter merges in train_step/4 go through the same function:

# before
model_state =
  update_in(model_state, [Access.key!(:data)], fn data ->
    tree_merge(data, trainable_parameters, fn _, _, v -> v end)
  end)

# after
model_state = Axon.ModelState.update(model_state, trainable_parameters)

The behaviour is the same. update/2 replaces the leaves of data that exist in trainable_parameters and keeps everything else; frozen parameters are never part of the gradient variable (trainable_parameters/1 diffs against frozen_parameters), so they stay untouched, and Axon.ModelState.SharedParameter leaves are skipped, so tied weights stay tied. updated_state defaults to %{}, which makes the state merge a no-op, and the parameters, state and frozen_parameters metadata are keep: fields of the container, so they survive the jitted step. Because the merge is a plain replacement, the same Nx.Defn.Expr tensors end up in data and gradients flow exactly as before. Layer state (batch-norm running statistics, etc.) is still produced by the forward pass and merged in the step function with the three-arity update/3 call, which is unchanged.

The only difference from the old private copy is that Axon.ModelState.tree_merge treats Axon.Quantization.QTensor as a leaf instead of recursing into the struct. This cannot change results because Axon.Quantization freezes quantized kernels, so they never appear in the trainable parameters.

Two tests were added to test/axon/loop_test.exs:

• train_step/3 updates only trainable parameters freezes one dense layer, runs a step, and checks that the frozen kernel and bias come back unchanged while the trainable layer's parameters move, and that parameters and frozen_parameters are preserved.
• train_step/3 updates nested parameters uses a reused Axon.block with a batch norm and checks that the nested block parameters and the layer state are both updated through the single path, with parameters and state preserved.

🤖 Generated with Claude Code