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,