View Source Axon.StatefulOutput (Axon v0.6.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".