View Source Axon.StatefulOutput (Axon v0.7.0)

Container for returning stateful outputs from Axon layers.

Some layers, such as Axon.batch_norm/2, keep a running internal state which is updated continuously at train time and used statically at inference time. In order for the Axon compiler to differentiate ordinary layer outputs from internal state, you must mark output as stateful.

Stateful Outputs consist of two fields:

:output - Actual layer output to be forwarded to next layer
:state - Internal layer state to be tracked and updated

:output is simply forwarded to the next layer. :state is aggregated with other stateful outputs, and then is treated specially by internal Axon training functions such that update state parameters reflect returned values from stateful outputs.

:state must be a map with keys that map directly to layer internal state names. For example, Axon.Layers.batch_norm returns StatefulOutput with :state keys of "mean" and "var".