diff options
| -rw-r--r-- | numpy/core/src/multiarray/einsum.c.src | 20 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/einsum_sumprod.c.src | 25 | ||||
| -rw-r--r-- | numpy/core/tests/test_einsum.py | 15 |
3 files changed, 36 insertions, 24 deletions
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src index 856dce11b..fe7fd0d32 100644 --- a/numpy/core/src/multiarray/einsum.c.src +++ b/numpy/core/src/multiarray/einsum.c.src @@ -534,11 +534,13 @@ unbuffered_loop_nop1_ndim2(NpyIter *iter) * loop provided by the iterator is in Fortran-order. */ int needs_api = NpyIter_IterationNeedsAPI(iter); - NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + if (!needs_api) { + NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + } for (coord = shape[1]; coord > 0; --coord) { sop(1, ptrs[0], strides[0], shape[0]); - if(needs_api && PyErr_Occurred()){ + if (needs_api && PyErr_Occurred()){ return -1; } @@ -593,12 +595,14 @@ unbuffered_loop_nop1_ndim3(NpyIter *iter) * loop provided by the iterator is in Fortran-order. */ int needs_api = NpyIter_IterationNeedsAPI(iter); - NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + if (!needs_api) { + NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + } for (coords[1] = shape[2]; coords[1] > 0; --coords[1]) { for (coords[0] = shape[1]; coords[0] > 0; --coords[0]) { sop(1, ptrs[0], strides[0], shape[0]); - if(needs_api && PyErr_Occurred()){ + if (needs_api && PyErr_Occurred()){ return -1; } @@ -655,7 +659,9 @@ unbuffered_loop_nop2_ndim2(NpyIter *iter) * loop provided by the iterator is in Fortran-order. */ int needs_api = NpyIter_IterationNeedsAPI(iter); - NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + if (!needs_api) { + NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + } for (coord = shape[1]; coord > 0; --coord) { sop(2, ptrs[0], strides[0], shape[0]); @@ -716,7 +722,9 @@ unbuffered_loop_nop2_ndim3(NpyIter *iter) * loop provided by the iterator is in Fortran-order. */ int needs_api = NpyIter_IterationNeedsAPI(iter); - NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + if (!needs_api) { + NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + } for (coords[1] = shape[2]; coords[1] > 0; --coords[1]) { for (coords[0] = shape[1]; coords[0] > 0; --coords[0]) { sop(2, ptrs[0], strides[0], shape[0]); diff --git a/numpy/core/src/multiarray/einsum_sumprod.c.src b/numpy/core/src/multiarray/einsum_sumprod.c.src index fa59223fe..5f2ccf2dc 100644 --- a/numpy/core/src/multiarray/einsum_sumprod.c.src +++ b/numpy/core/src/multiarray/einsum_sumprod.c.src @@ -1032,20 +1032,14 @@ bool_sum_of_products_outstride0_@noplabel@(int nop, char **dataptr, * object_sum_of_products_one, * object_sum_of_products_two, * object_sum_of_products_three, + * object_sum_of_products_contig_any, * object_sum_of_products_contig_one, - * object_sum_of_products_contig_two, - * object_sum_of_products_stride0_contig_outcontig_two, - * object_sum_of_products_contig_stride0_outcontig_two, - * object_sum_of_products_stride0_contig_outstride0_two, - * object_sum_of_products_contig_stride0_outstride0_two, - * object_sum_of_products_contig_contig_outstride0_two, + * object_sum_of_products_contig_two, * object_sum_of_products_contig_three, - * object_sum_of_products_contig_any, - * object_sum_of_products_contig_outstride0_one, + * object_sum_of_products_outstride0_any, * object_sum_of_products_outstride0_one, * object_sum_of_products_outstride0_two, - * object_sum_of_products_outstride0_three, - * object_sum_of_products_outstride0_any# + * object_sum_of_products_outstride0_three# */ static void @fn_name@(int nop, char **dataptr, @@ -1062,17 +1056,12 @@ static void if (!curr) { curr = Py_None; // convention is to treat nulls as None } - Py_SETREF(prod, PyNumber_Multiply(curr, prod)); + Py_SETREF(prod, PyNumber_Multiply(prod, curr)); if (!prod) { return; } } - if (*(PyObject **)dataptr[nop] == NULL) { - Py_DECREF(prod); - return; - } - PyObject *sum = PyNumber_Add(*(PyObject **)dataptr[nop], prod); Py_DECREF(prod); if (!sum) { @@ -1110,7 +1099,7 @@ _contig_outstride0_unary_specialization_table[NPY_NTYPES] = { * 1, 1, * 1, 1, 1, * 1, 1, 1, - * 1, 0, 0, 0, + * 0, 0, 0, 0, * 0, 0, 1# */ #if @use@ @@ -1141,7 +1130,7 @@ static sum_of_products_fn _binary_specialization_table[NPY_NTYPES][5] = { * 1, 1, * 1, 1, 1, * 0, 0, 0, - * 1, 0, 0, 0, + * 0, 0, 0, 0, * 0, 0, 1# */ #if @use@ diff --git a/numpy/core/tests/test_einsum.py b/numpy/core/tests/test_einsum.py index bcf23890c..3a06d119f 100644 --- a/numpy/core/tests/test_einsum.py +++ b/numpy/core/tests/test_einsum.py @@ -699,6 +699,21 @@ class TestEinsum: # see issue gh-15776 and issue gh-15256 assert_equal(np.einsum('i,j', [1], [2], out=None), [[2]]) + def test_object_loop(self): + + class Mult: + def __mul__(self, other): + return 42 + + objMult = np.array([Mult()]) + objNULL = np.ndarray(buffer = b'\0' * np.intp(0).itemsize, shape=1, dtype=object) + + with pytest.raises(TypeError): + np.einsum("i,j", [1], objNULL) + with pytest.raises(TypeError): + np.einsum("i,j", objNULL, [1]) + assert np.einsum("i,j", objMult, objMult) == 42 + def test_subscript_range(self): # Issue #7741, make sure that all letters of Latin alphabet (both uppercase & lowercase) can be used # when creating a subscript from arrays |