Add Axon.Loop.stream/4 and implement run/4 on top of it

[seanmor5]

Training loops are a map-reduce over a dataset, but Axon.Loop.run/4 collapses that into a single output. Any per-epoch logic — evaluating on a validation set, adjusting the dataset, logging something custom — has to be smuggled into a handle_event(:epoch_completed, ...) handler that threads its results through handler_metadata or a closure.

Axon.Loop.stream/4 returns a lazy stream which emits the accumulated %Axon.Loop.State{} at the end of every completed epoch, so per-epoch logic becomes ordinary Enum/Stream code:

loop
|> Axon.Loop.stream(train_data)
|> Stream.map(fn state ->
  {state.epoch, evaluate(state.step_state.model_state), state.metrics}
end)
|> Enum.take(10)

The stream is infinite by default

Rather than declaring an epoch count up front, the consumer decides how long to train:

loop |> Axon.Loop.stream(data) |> Enum.take(10)             # 10 epochs
loop |> Axon.Loop.stream(data) |> Stream.take_while(...)     # until it stops improving

:epochs is still accepted if you want a bounded stream. The tradeoff is that Axon.Loop.stream(loop, data) |> Enum.to_list() runs forever — inherent to an infinite stream, and called out in the docs.

run/4 is now a reduce over the stream

Both functions share one engine, a private loop_stream/5 built on Stream.resource/3 which drives the loop one epoch per next_fun call and emits tagged elements — {:epoch, state} per completed epoch, {:done, state} once at the end with the epoch => metrics map and :status.

def run(loop, data, init_state \\ %{}, opts \\ []) do
  %Loop{output_transform: output_transform} = loop

  loop
  |> loop_stream(data, init_state, opts, 1)
  |> Enum.reduce(nil, fn
    {:epoch, _state}, acc -> acc
    {:done, state}, _acc -> state
  end)
  |> output_transform.()
end

The epoch body is a direct transcription of the old Enum.reduce_while over epoch_start..epoch_end//1, so event ordering, caching of the compiled batch function across epochs, halt semantics and metric zero-filling are unchanged. run/4 keeps epochs: 1; changing that would break every existing caller, and "reduce to a single value" has no natural infinite default.

Semantics worth reviewing

• Emits the full %State{}, not output_transform.(state). Otherwise a trainer's stream would yield only model state and drop the metrics, defeating the point. loop.output_transform.(state) is available if you want it.
• :metrics on an emitted state is that epoch's metrics, not the epoch => metrics map run/4 returns. :step_state and :times are cumulative.
• Halted epochs emit nothing; :halt_loop ends the stream. This keeps "one element per completed epoch" crisp — the alternative would emit duplicate :epoch values, since the existing :epoch_halted/:continue path deliberately doesn't advance state.epoch.
• %State{max_epoch: :infinity} is now possible. The field is write-only in the codebase — nothing reads it — so this is doc-only fallout.
• Re-enumerating the stream re-initializes and re-runs the whole loop, per normal Stream semantics.

Tests

835 pass (was 826). New describe "streaming" block covers per-epoch emission (epoch/iteration/step_state/times), epochs: 0, unbounded streams driven three ways off one stream value, laziness verified with an Agent-counting data stream, re-enumerability, :halt_loop terminating an unbounded stream, halted-epoch suppression, per-epoch metrics, and run/4 agreeing with the final stream state for a seeded trainer.

🤖 Generated with Claude Code