Struct BatchNorm 

pub struct BatchNorm {
    pub weight: Parameter,
    pub bias: Parameter,
    /* private fields */
}

Batch normalization over the first (batch) dimension.

Uses a single fused torch::batch_norm kernel (1 autograd node). Training mode: uses batch statistics, updates running stats in-place. Eval mode: uses running statistics (persisted via buffers()).

Fields

weight: Parameter

bias: Parameter

Implementations

impl BatchNorm

pub fn new(num_features: i64) -> Result<Self>

Create a BatchNorm layer for num_features channels/features on CPU.

Examples found in repository

examples/showcase/main.rs (line 77)

fn silu_block(dim: i64) -> flodl::Result<Graph> {
    FlowBuilder::from(Linear::new(dim, dim)?)
        .through(SiLU)
        .through(BatchNorm::new(dim)?)
        .build()
}

pub fn on_device(num_features: i64, device: Device) -> Result<Self>

Create a BatchNorm layer on a specific device.

Trait Implementations

impl Module for BatchNorm

fn name(&self) -> &str

Human-readable type name used as node ID prefix in graph visualization. Override to return a lowercase identifier (e.g., “linear”, “gelu”).

fn forward(&self, input: &Variable) -> Result<Variable>

Run the forward pass on input and return the result.

fn parameters(&self) -> Vec<Parameter>

Return this module’s learnable parameters. Default: recursively collects from sub_modules() with pointer dedup. Leaf modules should override to return their own parameters.

fn buffers(&self) -> Vec<Buffer>

Return this module’s non-learnable persistent buffers (e.g., running stats). Default: recursively collects from sub_modules() with pointer dedup. Leaf modules should override to return their own buffers.

fn set_training(&self, training: bool)

Set training/eval mode. Affects Dropout, BatchNorm, etc. Override in modules with mode-dependent behavior.

fn move_to_device(&self, device: Device)

Move all parameters and buffers to the given device. Override in modules like BatchNorm that hold non-parameter state.

fn sub_modules(&self) -> Vec<Rc<dyn Module>>

Return direct child modules for recursive tree walks. Override in composite modules (loops, switches, gates).

fn train(&self)

Set training mode. Shorthand for set_training(true).

fn eval(&self)

Set eval mode. Shorthand for set_training(false).

fn trace(&self) -> Option<Variable>

Return per-iteration side output for loop tracing. Override in loop body modules that capture trajectory data (e.g., attention fixation points). Returns None by default. When Some, the loop executor collects traces accessible via Graph::traces().

fn as_named_input(&self) -> Option<&dyn NamedInputModule>

Upcast to NamedInputModule for multi-input graphs. Override in types that implement NamedInputModule to enable receiving additional named inputs via graph using().

fn structural_hash(&self) -> Option<String>

SHA-256 hex hash of module architecture for checkpoint validation. Override in composite modules (Graph) that compute a deterministic hash from their topology and parameter shapes.

fn reset(&self)

Reset internal state (e.g. recurrent hidden state) between sequences. Called by loops before iterating to clear stale tensors whose grad_fns may reference freed saved tensors. Override in stateful modules.

fn detach_state(&self)

Detach internal state from the computation graph (for truncated BPTT). Called between training steps to break gradient chains on state carried across forward passes (e.g., recurrent hidden state). Override in stateful modules.