Function layernorm 

pub fn layernorm(
    x: &WgpuStorage,
    gamma: &WgpuStorage,
    beta: &WgpuStorage,
    rows: usize,
    cols: usize,
    eps: f32,
) -> WgpuStorage