Module `Batch_norm.Stats`

Running statistics: the exponential moving averages of per-feature batch mean and (population) variance maintained by training-mode apply and consumed by eval-mode apply. Not parameters — never differentiate or optimize them; thread them through the training loop as auxiliary state.

type 'b stats = {

mean : (float, 'b) Nx.t;
var : (float, 'b) Nx.t;

}

The type for running statistics with float dtype layout 'b, each of shape [| features |].

type t = Nx.float32_elt stats

The type for single-precision running statistics, the common case. The exponential moving average accumulates small increments, so keep statistics at float32 even when the parameters are half precision.

val map : ('a 'c. ('a, 'c) Nx.t -> ('a, 'c) Nx.t) -> 'b stats -> 'b stats

map f s is s with f applied to mean and var.

val map2 : 
  ('a 'c. ('a, 'c) Nx.t -> ('a, 'c) Nx.t -> ('a, 'c) Nx.t) ->
  'b stats ->
  'b stats ->
  'b stats

map2 f s s' combines s and s' leafwise with f.

val iter : ('a 'c. ('a, 'c) Nx.t -> unit) -> 'b stats -> unit

iter f s applies f to mean and var, in that order.

val astype : (float, 'c) Nx.dtype -> 'b stats -> 'c stats

astype dt s is s with mean and var cast to dt.

val names : 'b stats -> string list

names s is ["mean"; "var"], pairing leaves in traversal order (see Checkpoint.Named).