Initialize training step state in the types the step produces

[seanmor5]

Closes #568.

The problem

Axon.Loop.run/4 compiles the batch function once, against the initial step state and metrics as templates, and calls that compiled function on every following batch. That only works if the step state is a fixed point of the step in shape and type. On main it is not as soon as anything in the pipeline is wider than f32. The Nx.while error from the original report is gone (that code no longer exists), but the same defect now surfaces on the second batch of any f64 run:

** (ArgumentError) argument at position 3 is not compatible with compiled function template.
%{i: s32, loss: <<<<< Expected f32 ==== Argument f64 >>>>>, ...

Three places hardcoded f32 where the step produces something else:

1. train_step/4 seeded loss with Nx.tensor(0.0); the step computes loss * i + batch_loss / (i + 1), so with an f64 loss it turns f64 after one step.
2. Polaris initializes optimizer moments in f32 regardless of the parameter type; with f64 gradients Adam's mu/nu promote to f64 on the first update. Every optimizer with moments (adam, adamw, lamb, radam, rmsprop, yogi, adabelief, sgd with momentum) does this.
3. run/4 started every metric accumulator from Nx.tensor(0, type: :f32) (next to a # TODO: Can we infer here?), so a "loss" metric copied from an f64 step state, or mean_absolute_error over f64 predictions, changed type after the first batch too.

The existing integration test "f64 input test" (f64 inputs against an f32 model) fails on main for the same reason and passes with this change.

The design

Rather than special-casing each entry, the step state is initialized in the types the step produces, learned by tracing one step, not computing it. train_step/4's init already traces the forward pass only to read the prediction's shape and type; this extends the same idea to the whole step. The step body is split out of step_fn so init can call it, and init casts its seed state to the traced result:

cast_like(state, step_body.({inp, tar}, state))

cast_like/2 only reads Nx.type/1 off the traced leaves, so nothing of the gradient expression is referenced by the returned state or lowered; the cost is one extra trace at init. One trace is enough because Nx type promotion is monotone: after one step every accumulator has already merged with the types of everything it merges with. Nx.as_type/2 is a no-op where the type already matches, so i, the loss-scale state and every already-correct tensor are untouched.

Metrics get the same treatment: init_metrics/4 traces one accumulation from an f32 zero over the initial step state and starts each accumulator from zeros_like of that value. Metrics that are read straight from the step state (the trainer's "loss") take the step state's type.

Nothing is guessed by hand: the loss type is whatever the loss function returns for the model's output and targets, and optimizer state is whatever the optimizer's update produces. That is what makes this hold for custom losses, custom optimizers and every loss-scale mode without per-case code.

Usage

policy = Axon.MixedPrecision.create_policy(params: {:f, 64}, compute: {:f, 64}, output: {:f, 64})

model =
  Axon.input("input", shape: {nil, 4})
  |> Axon.dense(8, activation: :relu)
  |> Axon.dense(1)
  |> Axon.MixedPrecision.apply_policy(policy)

model
|> Axon.Loop.trainer(:mean_squared_error, :adam)
|> Axon.Loop.metric(:mean_absolute_error)
|> Axon.Loop.run(data, Axon.ModelState.empty(), epochs: 5)

This now runs to completion with the default strict?: true; the loss, the Adam moments and both metrics are f64.

Tradeoffs and limitations

• f16/bf16 behaviour is unchanged. With f16 parameters the gradients and loss are f32 and Polaris keeps its f32 master state, so the state was already a fixed point; the cast finds the same types and downcasts nothing. A test pins this.
• Metrics still start from an f32 zero, so for f32 models nothing changes: integer-valued metrics accumulated with :running_sum keep accumulating in f32 as they do today. Only metrics whose values are wider than f32 widen.
• Init now traces the loss and optimizer update, so the targets given to init_fn must be valid for the loss. run/4 always initializes from a real batch so this is transparent there; two tests that called init_fn directly with placeholder targets of the wrong shape have been updated to pass real ones.
• Polaris's f32 moment initialization is left alone; working around it explicitly would break the f16 master-state case and would still leave custom optimizers unhandled.

Tests

• train_step/3 with an f64 policy: loss and Adam moments are f64 at init, and the leaf types of the init state equal those of the state after one step, for :identity, :dynamic and :static loss scales. Fails on main with loss and eight Adam tensors differing.
• train_step/3 with an f16 policy keeps f32 loss and moments, f16 parameters, and is a type fixed point (guards against downcasting).
• A strict Axon.Loop.run of an f64 policy model with :adam and a mean_absolute_error metric completes two epochs with f64 loss, parameters and metrics. Raises the template error on main.
• A strict run with f64 targets against an f32 model completes with :sgd and :adam and reports an f64 loss (the same failure mode as the existing "f64 input test" integration test, which now passes).

The full suite passes, test/axon/loop_test.exs passes on EXLA, and test/axon/integration_test.exs --include integration passes apart from "image classification test", which times out at 60s on main as well on the default backend.

🤖 Generated with Claude Code