Module `Kaun.Layer_norm`

Layer normalization over the feature axis.

Layer normalization (Ba, Kiros and Hinton, 2016).

Layer norm standardizes each feature vector — the last axis of the input — to zero mean and unit variance, then applies a learned per-feature scale (gamma) and shift (beta). Unlike batch normalization it is stateless and independent of the batch: the same function at training and inference time. Construct parameters with init or make and normalize with apply; map, map2, iter and names supply the Nx.Ptree.S and checkpoint plumbing.

Types

type 'b params = {

gamma : (float, 'b) Nx.t;
beta : (float, 'b) Nx.t;

}

The type for layer-norm parameters with float dtype layout 'b. gamma (the scale) and beta (the shift) both have shape [| dim |], one entry per normalized feature.

type t = Nx.float32_elt params

The type for single-precision layer norms, the common case.

Constructors

val make : dim:int -> (float, 'b) Nx.dtype -> 'b params

make ~dim dtype is a fresh identity normalization over dim features: gamma all ones, beta all zeros.

Raises Invalid_argument if dim is not positive.

val init : dim:int -> t

init ~dim is make ~dim Nx.float32.

Applying

val apply : ?eps:float -> 'b params -> (float, 'b) Nx.t -> (float, 'b) Nx.t

apply p x normalizes each vector along x's last axis and rescales it:

(x - mean(x)) / sqrt (var(x) + eps) * gamma + beta

where mean and var (the biased variance) are taken along the last axis; every other axis is treated as a batch axis. eps keeps the division finite for constant vectors — a constant vector maps to beta — and defaults to 1e-5. The result has x's shape. Differentiable through Rune.

For half and quarter precision inputs (float16, bfloat16, float8) the statistics and normalization are computed in a float32 island: x is upcast, normalized at float32, and the normalized values are cast back to x's dtype before the gamma/beta affine transform. Float32 and float64 inputs use their own dtype throughout, exactly as if the island were absent.

Raises Invalid_argument if x is a scalar, if x's last axis does not have size dim, or if eps is negative.

Traversals

Plain traversals over the parameter leaves, in the order gamma then beta. They satisfy the Nx.Ptree.S contract at any fixed 'b.

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

map f p is p with f applied to every parameter leaf.

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

map2 f p q combines p and q leafwise with f.

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

iter f p applies f to every parameter leaf of p.

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

astype dt p is p with every parameter leaf cast to dt. Differentiable through Rune: gradients flow back at each original leaf's dtype, so an astype of float32 parameters inside a loss function yields float32 gradients.

val names : 'b params -> string list

names p is ["gamma"; "beta"], the checkpoint names of the parameter leaves in traversal order. See Checkpoint.Named.