Softmax numerical stability. (#267)

* Softmax numerical stability.

* Fix the flash-attn test.

candle-nn/tests/ops.rs

@@ -0,0 +1,62 @@
+use candle::{Device, Result, Tensor};
+
+pub fn to_vec3_round(t: Tensor, digits: i32) -> Result<Vec<Vec<Vec<f32>>>> {
+    let b = 10f32.powi(digits);
+    let t = t.to_vec3::<f32>()?;
+    let t = t
+        .iter()
+        .map(|t| {
+            t.iter()
+                .map(|t| t.iter().map(|t| f32::round(t * b) / b).collect())
+                .collect()
+        })
+        .collect();
+    Ok(t)
+}
+
+#[test]
+fn softmax() -> Result<()> {
+    let device = &Device::Cpu;
+    let data = &[[[3f32, 1., 4.], [1., 5., 9.]], [[2., 1., 7.], [8., 2., 8.]]];
+    let tensor = Tensor::new(data, device)?;
+    let t0 = candle_nn::ops::softmax(&tensor.log()?, 0)?;
+    let t1 = candle_nn::ops::softmax(&tensor.log()?, 1)?;
+    let t2 = candle_nn::ops::softmax(&tensor.log()?, 2)?;
+    assert_eq!(
+        to_vec3_round(t0, 4)?,
+        &[
+            // 3/5, 1/2, 4/11
+            [[0.6, 0.5, 0.3636], [0.1111, 0.7143, 0.5294]],
+            // 2/5, 1/2, 7/11
+            [[0.4, 0.5, 0.6364], [0.8889, 0.2857, 0.4706]]
+        ]
+    );
+    assert_eq!(
+        to_vec3_round(t1, 4)?,
+        &[
+            // 3/4, 1/6, 4/13
+            [[0.75, 0.1667, 0.3077], [0.25, 0.8333, 0.6923]],
+            // 2/10, 1/3, 7/15
+            [[0.2, 0.3333, 0.4667], [0.8, 0.6667, 0.5333]]
+        ]
+    );
+    assert_eq!(
+        to_vec3_round(t2, 4)?,
+        &[
+            // (3, 1, 4) / 8, (1, 5, 9) / 15
+            [[0.375, 0.125, 0.5], [0.0667, 0.3333, 0.6]],
+            // (2, 1, 7) / 10, (8, 2, 8) / 18
+            [[0.2, 0.1, 0.7], [0.4444, 0.1111, 0.4444]]
+        ]
+    );
+    Ok(())
+}
+
+#[test]
+fn softmax_numerical_stability() -> Result<()> {
+    let dev = &Device::Cpu;
+    let xs = Tensor::new(&[1234f32, 0.], dev)?;
+    let softmax = candle_nn::ops::softmax(&xs, 0)?;
+    assert_eq!(softmax.to_vec1::<f32>()?, &[1f32, 0.]);
+    Ok(())
+}