Function burn_layer_norm_nd 

pub fn burn_layer_norm_nd<const D: usize>(
    x: Tensor<D>,
    w: Tensor<1>,
    b: Tensor<1>,
    eps: f32,
) -> Tensor<D>

LayerNorm over the last dim for rank-D input.