Module norm 

Functions

burn_layer_norm

LayerNorm over the last dim, with weight and bias.

burn_layer_norm_nd

LayerNorm over the last dim for rank-D input.

burn_layer_norm_no_affine

LayerNorm without affine parameters over the last dim (2D).

burn_rms_norm

RMSNorm over the last dim, with weight.

burn_rms_norm_nd

RMSNorm over the last dim for rank-D input.

burn_softmax_last_dim

Softmax over the last dim for host-checked logits.