Option for precise softmax (#953)

* precise softmax

* Add an equivalency check

* Make the threadgroup memory definition fixed

* precise cpu softmax

* precise option on cpu

* remove print

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

python/src/ops.cpp

@@ -2430,12 +2430,13 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "softmax",
-      [](const array& a, const IntOrVec& axis, StreamOrDevice s) {
-        return softmax(a, get_reduce_axes(axis, a.ndim()), s);
+      [](const array& a, const IntOrVec& axis, bool precise, StreamOrDevice s) {
+        return softmax(a, get_reduce_axes(axis, a.ndim()), precise, s);
       },
       nb::arg(),
       "axis"_a = nb::none(),
       nb::kw_only(),
+      "precise"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
           "def softmax(a: array, /, axis: Union[None, int, Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"),

python/tests/test_ops.py

@@ -1430,6 +1430,13 @@ class TestOps(mlx_tests.MLXTestCase):
         out = mx.softmax(y[:, 0:2], axis=-1)
         self.assertAlmostEqual(out.sum().item(), 8.0, 5)
 
+        # Precise
+        for t in [mx.float16, mx.bfloat16]:
+            a = (10 * mx.random.normal(shape=(1024,))).astype(t)
+            out_expect = mx.softmax(a.astype(mx.float32)).astype(t)
+            out = mx.softmax(a, axis=-1, precise=True)
+            self.assertTrue(mx.allclose(out_expect, out))
+
     def test_concatenate(self):
         a_npy = np.random.randn(32, 32, 32)
         b_npy = np.random.randn(32, 32, 32)