https://github.com/scikit-hep/vector/issues/616 https://github.com/numpy/numpy/issues/29388 https://github.com/numpy/numpy/pull/29392 From d8eeb76d457a03347387c57ae7b84530aa38578d Mon Sep 17 00:00:00 2001 From: Maanas Arora Date: Wed, 16 Jul 2025 18:11:31 -0400 Subject: [PATCH 1/5] BUG: update fast_scalar_power to handle special-case squaring for any array type except object arrays --- numpy/_core/src/multiarray/number.c | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/numpy/_core/src/multiarray/number.c b/numpy/_core/src/multiarray/number.c index b801d7e041e2..8ab1d0d4a784 100644 --- a/numpy/_core/src/multiarray/number.c +++ b/numpy/_core/src/multiarray/number.c @@ -363,7 +363,12 @@ fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) } PyArrayObject *a1 = (PyArrayObject *)o1; - if (!(PyArray_ISFLOAT(a1) || PyArray_ISCOMPLEX(a1))) { + if (PyArray_ISOBJECT(a1)) { + return 1; + } + if (!is_square && !PyArray_ISFLOAT(a1) && !PyArray_ISCOMPLEX(a1)) { + // we special-case squaring for any array type + // gh-29388 return 1; } From 097b04a300e9bd5f35230a6450dbf7d939426dd3 Mon Sep 17 00:00:00 2001 From: Maanas Arora Date: Wed, 16 Jul 2025 18:19:04 -0400 Subject: [PATCH 2/5] BUG: fix missing declaration --- numpy/_core/src/multiarray/number.c | 3 +++ 1 file changed, 3 insertions(+) diff --git a/numpy/_core/src/multiarray/number.c b/numpy/_core/src/multiarray/number.c index 8ab1d0d4a784..6b885ffa68e1 100644 --- a/numpy/_core/src/multiarray/number.c +++ b/numpy/_core/src/multiarray/number.c @@ -332,6 +332,8 @@ static int fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) { PyObject *fastop = NULL; + int is_square = 0; + if (PyLong_CheckExact(o2)) { int overflow = 0; long exp = PyLong_AsLongAndOverflow(o2, &overflow); @@ -344,6 +346,7 @@ fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) } else if (exp == 2) { fastop = n_ops.square; + is_square = 1; } else { return 1; From 79a91a1bdaa50339ad513998c440b8949f15f56e Mon Sep 17 00:00:00 2001 From: Maanas Arora Date: Wed, 16 Jul 2025 18:29:08 -0400 Subject: [PATCH 3/5] TST: add test to ensure `arr**2` calls square for structured dtypes --- numpy/_core/tests/test_multiarray.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py index 04222025883e..364a9793a81f 100644 --- a/numpy/_core/tests/test_multiarray.py +++ b/numpy/_core/tests/test_multiarray.py @@ -4215,6 +4215,13 @@ def pow_for(exp, arr): assert_equal(obj_arr ** -1, pow_for(-1, obj_arr)) assert_equal(obj_arr ** 2, pow_for(2, obj_arr)) + def test_pow_calls_square_structured_dtype(self): + # gh-29388 + dt = np.dtype([('a', 'i4'), ('b', 'i4')]) + a = np.array([(1, 2), (3, 4)], dtype= dt) + with pytest.raises(TypeError, match="ufunc 'square' not supported"): + a ** 2 + def test_pos_array_ufunc_override(self): class A(np.ndarray): def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): From e014ffe34ee4573538691cd197273ff08f35d26b Mon Sep 17 00:00:00 2001 From: Maanas Arora Date: Wed, 16 Jul 2025 18:32:48 -0400 Subject: [PATCH 4/5] STY: remove whitespace --- numpy/_core/tests/test_multiarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py index 364a9793a81f..e7647d2e23fe 100644 --- a/numpy/_core/tests/test_multiarray.py +++ b/numpy/_core/tests/test_multiarray.py @@ -4218,7 +4218,7 @@ def pow_for(exp, arr): def test_pow_calls_square_structured_dtype(self): # gh-29388 dt = np.dtype([('a', 'i4'), ('b', 'i4')]) - a = np.array([(1, 2), (3, 4)], dtype= dt) + a = np.array([(1, 2), (3, 4)], dtype=dt) with pytest.raises(TypeError, match="ufunc 'square' not supported"): a ** 2 From ddaef11f219990372f1c10bca2d3d3f0bb57a2e9 Mon Sep 17 00:00:00 2001 From: Maanas Arora Date: Thu, 17 Jul 2025 05:16:12 -0400 Subject: [PATCH 5/5] BUG: replace new variable `is_square` with direct op comparison in `fast_scalar_power` function --- numpy/_core/src/multiarray/number.c | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/numpy/_core/src/multiarray/number.c b/numpy/_core/src/multiarray/number.c index 6b885ffa68e1..de4012641684 100644 --- a/numpy/_core/src/multiarray/number.c +++ b/numpy/_core/src/multiarray/number.c @@ -332,7 +332,6 @@ static int fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) { PyObject *fastop = NULL; - int is_square = 0; if (PyLong_CheckExact(o2)) { int overflow = 0; @@ -346,7 +345,6 @@ fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) } else if (exp == 2) { fastop = n_ops.square; - is_square = 1; } else { return 1; @@ -369,7 +367,7 @@ fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result) if (PyArray_ISOBJECT(a1)) { return 1; } - if (!is_square && !PyArray_ISFLOAT(a1) && !PyArray_ISCOMPLEX(a1)) { + if (fastop != n_ops.square && !PyArray_ISFLOAT(a1) && !PyArray_ISCOMPLEX(a1)) { // we special-case squaring for any array type // gh-29388 return 1;