From 9592766939f0be4eaba50965b63967634a701958 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Sun, 1 Sep 2024 19:49:16 -0700
Subject: [PATCH] add std as method (#1387)

* add std as method

* add std as method
---
 docs/src/python/array.rst |  3 ++-
 docs/src/python/ops.rst   |  1 +
 python/src/array.cpp      | 16 ++++++++++++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/docs/src/python/array.rst b/docs/src/python/array.rst
index 440167ba43..532bb45c9f 100644
--- a/docs/src/python/array.rst
+++ b/docs/src/python/array.rst
@@ -53,8 +53,9 @@ Array
     array.sqrt
     array.square
     array.squeeze
-    array.swapaxes
+    array.std
     array.sum
+    array.swapaxes
     array.transpose
     array.T
     array.var
diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst
index 3c58e0de50..c2e70c8244 100644
--- a/docs/src/python/ops.rst
+++ b/docs/src/python/ops.rst
@@ -78,6 +78,7 @@ Operations
    hadamard_transform
    identity
    inner
+   isfinite
    isclose
    isinf
    isnan
diff --git a/python/src/array.cpp b/python/src/array.cpp
index caf7ce6df6..71ecd658fc 100644
--- a/python/src/array.cpp
+++ b/python/src/array.cpp
@@ -1115,6 +1115,22 @@ void init_array(nb::module_& m) {
           nb::kw_only(),
           "stream"_a = nb::none(),
           "See :func:`mean`.")
+      .def(
+          "std",
+          [](const array& a,
+             const IntOrVec& axis,
+             bool keepdims,
+             int ddof,
+             StreamOrDevice s) {
+            return mlx::core::std(
+                a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s);
+          },
+          "axis"_a = nb::none(),
+          "keepdims"_a = false,
+          "ddof"_a = 0,
+          nb::kw_only(),
+          "stream"_a = nb::none(),
+          "See :func:`std`.")
       .def(
           "var",
           [](const array& a,