diff options
Diffstat (limited to 'candle-nn/tests/layer_norm.rs')
| -rw-r--r-- | candle-nn/tests/layer_norm.rs | 15 |
1 files changed, 12 insertions, 3 deletions
diff --git a/candle-nn/tests/layer_norm.rs b/candle-nn/tests/layer_norm.rs index f81c29bd..30f598b3 100644 --- a/candle-nn/tests/layer_norm.rs +++ b/candle-nn/tests/layer_norm.rs @@ -13,6 +13,12 @@ fn layer_norm() -> Result<()> { let device = &Device::Cpu; let w = Tensor::new(&[3f32], device)?; let b = Tensor::new(&[0.5f32], device)?; + let ln2 = LayerNorm::new(Tensor::cat(&[&w, &w], 0)?, Tensor::cat(&[&b, &b], 0)?, 1e-8); + let ln3 = LayerNorm::new( + Tensor::cat(&[&w, &w, &w], 0)?, + Tensor::cat(&[&b, &b, &b], 0)?, + 1e-8, + ); let ln = LayerNorm::new(w, b, 1e-8); let two = Tensor::new(&[[[2f32]]], device)?; @@ -20,11 +26,11 @@ fn layer_norm() -> Result<()> { assert_eq!(res.to_vec1::<f32>()?, [0.5f32]); let inp = Tensor::new(&[[[4f32, 0f32]]], device)?; - let res = ln.forward(&inp)?; + let res = ln2.forward(&inp)?; assert_eq!(res.to_vec3::<f32>()?, [[[3.5f32, -2.5]]]); let inp = Tensor::new(&[[[1f32, 2., 3.], [4., 5., 6.], [9., 8., 7.]]], device)?; - let res = ln.forward(&inp)?; + let res = ln3.forward(&inp)?; assert_eq!( test_utils::to_vec3_round(&res, 4)?, [[ @@ -35,7 +41,10 @@ fn layer_norm() -> Result<()> { ); let mean = (res.sum_keepdim(2)? / 3.0)?; // The average value should be `b`. - assert_eq!(mean.to_vec3::<f32>()?, [[[0.5], [0.5], [0.5]]]); + assert_eq!( + test_utils::to_vec3_round(&mean, 4)?, + [[[0.5], [0.5], [0.5]]] + ); let std = (res.broadcast_sub(&mean)?.sqr()?.sum_keepdim(2)?.sqrt()? / 3.0)?; // The standard deviation should be sqrt(`w`). assert_eq!( |