diffusion_rs_common::nn

Module layer_norm

source
Expand description

Layer Normalization.

This layer applies Layer Normalization over a mini-batch of inputs as described in Layer Normalization. The input is expected to have three dimensions: a batch dimension, a length, and a hidden size, the normalization is applied over the last dimension.

§Example

use diffusion_rs_common::core::{Tensor, Device::Cpu, test_utils::to_vec3_round};
use diffusion_rs_common::nn::{LayerNorm, Module};

let w = Tensor::new(&[1f32, 1f32, 1f32], &Cpu)?;
let b = Tensor::new(&[0f32, 0f32, 0f32], &Cpu)?;
let layer = LayerNorm::new(w, b, 1e-5);

let xs = Tensor::new(
    &[[[1f32, 2., 3.], [4., 5., 6.], [9., 8., 7.]]],
    &Cpu)?;
let ys = layer.forward(&xs)?;
assert_eq!(
    to_vec3_round(&ys, 4)?,
    &[[[-1.2247, 0.0,  1.2247],
       [-1.2247, 0.0,  1.2247],
       [ 1.2247, 0.0, -1.2247]]]);

Structs§

Functions§