use super::*;
use crate::{shapes::*, tensor::*};
pub trait StddevTo<E: Dtype>: HasErr + HasShape {
fn stddev<Dst: Shape, Ax: Axes>(self, epsilon: impl Into<f64>) -> Self::WithShape<Dst>
where
Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>,
{
self.try_stddev(epsilon).unwrap()
}
fn try_stddev<Dst: Shape, Ax: Axes>(
self,
epsilon: impl Into<f64>,
) -> Result<Self::WithShape<Dst>, Self::Err>
where
Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>;
}
impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> StddevTo<E> for Tensor<S, E, D, T> {
fn try_stddev<Dst: Shape, Ax: Axes>(
self,
epsilon: impl Into<f64>,
) -> Result<Self::WithShape<Dst>, Self::Err>
where
Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>,
{
self.try_var()?
.try_add(E::from_f64(epsilon.into()).unwrap())?
.try_sqrt()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tests::*;
#[test]
fn test_std_axis_0_2d() {
let dev: TestDevice = Default::default();
let t = dev
.tensor([[1.0, 2.0, 3.0, 4.0], [0.0, 2.0, 5.0, 10.0]])
.to_dtype::<TestDtype>();
let r = t.leaky_trace().stddev::<Rank1<4>, _>(1e-8);
assert_close_to_literal!(r, [0.5, 0.0001, 1.0, 3.0]);
let g = r.mean().backward();
assert_close_to_literal!(
g.get(&t),
[[0.125, 0.0, -0.125, -0.125], [-0.125, 0.0, 0.125, 0.125]]
);
}
#[test]
fn test_std_axis_1_2d() {
let dev: TestDevice = Default::default();
let t = dev
.tensor([[1.0, 2.0, 3.0, 4.0], [0.0, 2.0, 5.0, 10.0]])
.to_dtype::<TestDtype>();
let r = t.leaky_trace().stddev::<Rank1<2>, _>(0.0);
assert_close_to_literal!(r, [1.118034, 3.7666297]);
let g = r.mean().backward();
assert_close_to_literal!(
g.get(&t),
[
[-0.16770509, -0.0559017, 0.0559017, 0.16770509],
[-0.14104122, -0.07466887, 0.024889633, 0.19082046],
]
);
}
}