diff --git a/python/src/array.cpp b/python/src/array.cpp index fbab846182..3d9e3f7d1d 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -1352,7 +1352,9 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::std( a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, @@ -1360,6 +1362,7 @@ void init_array(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), "See :func:`std`.") .def( @@ -1368,7 +1371,9 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::var( a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, @@ -1376,6 +1381,7 @@ void init_array(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), "See :func:`var`.") .def( diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 677d998af1..8ecc3122f7 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -2856,7 +2856,9 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), @@ -2864,6 +2866,7 @@ void init_ops(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), nb::sig( "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), @@ -2889,7 +2892,9 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + ddof = correction.value_or(ddof); return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), @@ -2897,6 +2902,7 @@ void init_ops(nb::module_& m) { "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), nb::sig( "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"),