From 8864647b654ecafbcf2ce27c551d30c6217f7161 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:05:06 +0530 Subject: [PATCH 1/2] Add a correction parameter in std and var Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- python/src/array.cpp | 20 ++++++++++++++++++-- python/src/ops.cpp | 28 ++++++++++++++++++++++++---- python/tests/test_ops.py | 8 ++++++++ 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/python/src/array.cpp b/python/src/array.cpp index fbab846182..935308694b 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -1352,14 +1352,22 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + int dof = correction ? static_cast(*correction) : ddof; + if (correction && ddof != 0 && + static_cast(*correction) != ddof) { + throw std::invalid_argument( + "[std] At most one of ddof or correction can be specified."); + } return mx::std( - a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); + a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); }, "axis"_a = nb::none(), "keepdims"_a = false, "ddof"_a = 0, nb::kw_only(), + "correction"_a = nb::none(), "stream"_a = nb::none(), "See :func:`std`.") .def( @@ -1368,14 +1376,22 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { + int dof = correction ? static_cast(*correction) : ddof; + if (correction && ddof != 0 && + static_cast(*correction) != ddof) { + throw std::invalid_argument( + "[var] At most one of ddof or correction can be specified."); + } return mx::var( - a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); + a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); }, "axis"_a = nb::none(), "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..946bb83e42 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -2856,17 +2856,24 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { - return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); + int dof = correction ? static_cast(*correction) : ddof; + if (correction && ddof != 0 && static_cast(*correction) != ddof) { + throw std::invalid_argument( + "[var] At most one of ddof or correction can be specified."); + } + return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); }, nb::arg(), "axis"_a = nb::none(), "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"), + "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, correction: float | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the variance(s) over the given axes. @@ -2879,6 +2886,9 @@ void init_ops(nb::module_& m) { singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. + correction (float, optional): Array API compatible alias for + ``ddof``. At most one of ``ddof`` or ``correction`` may be + given a non-default value. Returns: array: The output array of variances. @@ -2889,17 +2899,24 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, + std::optional correction, mx::StreamOrDevice s) { - return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); + int dof = correction ? static_cast(*correction) : ddof; + if (correction && ddof != 0 && static_cast(*correction) != ddof) { + throw std::invalid_argument( + "[std] At most one of ddof or correction can be specified."); + } + return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); }, nb::arg(), "axis"_a = nb::none(), "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"), + "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, correction: float | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the standard deviation(s) over the given axes. @@ -2912,6 +2929,9 @@ void init_ops(nb::module_& m) { singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. + correction (float, optional): Array API compatible alias for + ``ddof``. At most one of ``ddof`` or ``correction`` may be + given a non-default value. Returns: array: The output array of standard deviations. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 0c94989e02..876db1bf95 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1065,6 +1065,10 @@ def test_var(self): self.assertEqual(mx.var(x).dtype, mx.float32) self.assertAlmostEqual(mx.var(x).item(), x_np.var().item(), places=5) + self.assertEqual(mx.var(x, correction=1.0).item(), mx.var(x, ddof=1).item()) + with self.assertRaises(ValueError): + mx.var(x, ddof=1, correction=2.0) + def test_std(self): x = mx.random.uniform(shape=(5, 5)) x_np = np.array(x) @@ -1075,6 +1079,10 @@ def test_std(self): self.assertEqual(mx.std(x).dtype, mx.float32) self.assertAlmostEqual(mx.std(x).item(), x_np.std().item(), places=5) + self.assertEqual(mx.std(x, correction=1.0).item(), mx.std(x, ddof=1).item()) + with self.assertRaises(ValueError): + mx.std(x, ddof=1, correction=2.0) + def test_abs(self): a = mx.array([-1.0, 1.0, -2.0, 3.0]) result = mx.abs(a) From 4957a798e7082b26cac4580c0bc0ccf00cc20d28 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 26 Aug 2026 09:24:50 +0900 Subject: [PATCH 2/2] Minimal compatibility --- python/src/array.cpp | 22 ++++++---------------- python/src/ops.cpp | 30 ++++++++---------------------- python/tests/test_ops.py | 8 -------- 3 files changed, 14 insertions(+), 46 deletions(-) diff --git a/python/src/array.cpp b/python/src/array.cpp index 935308694b..3d9e3f7d1d 100644 --- a/python/src/array.cpp +++ b/python/src/array.cpp @@ -1352,16 +1352,11 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, - std::optional correction, + std::optional correction, mx::StreamOrDevice s) { - int dof = correction ? static_cast(*correction) : ddof; - if (correction && ddof != 0 && - static_cast(*correction) != ddof) { - throw std::invalid_argument( - "[std] At most one of ddof or correction can be specified."); - } + ddof = correction.value_or(ddof); return mx::std( - a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); + a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, "axis"_a = nb::none(), "keepdims"_a = false, @@ -1376,16 +1371,11 @@ void init_array(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, - std::optional correction, + std::optional correction, mx::StreamOrDevice s) { - int dof = correction ? static_cast(*correction) : ddof; - if (correction && ddof != 0 && - static_cast(*correction) != ddof) { - throw std::invalid_argument( - "[var] At most one of ddof or correction can be specified."); - } + ddof = correction.value_or(ddof); return mx::var( - a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); + a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, "axis"_a = nb::none(), "keepdims"_a = false, diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 946bb83e42..8ecc3122f7 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -2856,14 +2856,10 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, - std::optional correction, + std::optional correction, mx::StreamOrDevice s) { - int dof = correction ? static_cast(*correction) : ddof; - if (correction && ddof != 0 && static_cast(*correction) != ddof) { - throw std::invalid_argument( - "[var] At most one of ddof or correction can be specified."); - } - return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); + ddof = correction.value_or(ddof); + return mx::var(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), "axis"_a = nb::none(), @@ -2873,7 +2869,7 @@ void init_ops(nb::module_& m) { "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, *, correction: float | None = None, stream: StreamOrDevice = None) -> array"), + "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the variance(s) over the given axes. @@ -2886,9 +2882,6 @@ void init_ops(nb::module_& m) { singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. - correction (float, optional): Array API compatible alias for - ``ddof``. At most one of ``ddof`` or ``correction`` may be - given a non-default value. Returns: array: The output array of variances. @@ -2899,14 +2892,10 @@ void init_ops(nb::module_& m) { const IntOrVec& axis, bool keepdims, int ddof, - std::optional correction, + std::optional correction, mx::StreamOrDevice s) { - int dof = correction ? static_cast(*correction) : ddof; - if (correction && ddof != 0 && static_cast(*correction) != ddof) { - throw std::invalid_argument( - "[std] At most one of ddof or correction can be specified."); - } - return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, dof, s); + ddof = correction.value_or(ddof); + return mx::std(a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); }, nb::arg(), "axis"_a = nb::none(), @@ -2916,7 +2905,7 @@ void init_ops(nb::module_& m) { "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, *, correction: float | None = None, stream: StreamOrDevice = None) -> array"), + "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the standard deviation(s) over the given axes. @@ -2929,9 +2918,6 @@ void init_ops(nb::module_& m) { singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. - correction (float, optional): Array API compatible alias for - ``ddof``. At most one of ``ddof`` or ``correction`` may be - given a non-default value. Returns: array: The output array of standard deviations. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 876db1bf95..0c94989e02 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1065,10 +1065,6 @@ def test_var(self): self.assertEqual(mx.var(x).dtype, mx.float32) self.assertAlmostEqual(mx.var(x).item(), x_np.var().item(), places=5) - self.assertEqual(mx.var(x, correction=1.0).item(), mx.var(x, ddof=1).item()) - with self.assertRaises(ValueError): - mx.var(x, ddof=1, correction=2.0) - def test_std(self): x = mx.random.uniform(shape=(5, 5)) x_np = np.array(x) @@ -1079,10 +1075,6 @@ def test_std(self): self.assertEqual(mx.std(x).dtype, mx.float32) self.assertAlmostEqual(mx.std(x).item(), x_np.std().item(), places=5) - self.assertEqual(mx.std(x, correction=1.0).item(), mx.std(x, ddof=1).item()) - with self.assertRaises(ValueError): - mx.std(x, ddof=1, correction=2.0) - def test_abs(self): a = mx.array([-1.0, 1.0, -2.0, 3.0]) result = mx.abs(a)