pub fn layer_norm_axis<A, S, D>( x: &ArrayBase<S, D>, axis: Axis, eps: f64, ) -> ArrayBase<OwnedRepr<A>, D>where A: Float + FromPrimitive, D: RemoveAxis, S: Data<Elem = A>,