diff options
Diffstat (limited to 'numpy/core/src/multiarray/einsum.c.src')
| -rw-r--r-- | numpy/core/src/multiarray/einsum.c.src | 42 |
1 files changed, 35 insertions, 7 deletions
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src index cd1a58982..0b3b3fd8c 100644 --- a/numpy/core/src/multiarray/einsum.c.src +++ b/numpy/core/src/multiarray/einsum.c.src @@ -17,6 +17,7 @@ #include <numpy/npy_common.h> #include <numpy/arrayobject.h> #include <npy_pycompat.h> +#include <array_assign.h> //PyArray_AssignRawScalar #include <ctype.h> @@ -321,8 +322,7 @@ get_single_op_view(PyArrayObject *op, char *labels, Py_TYPE(op), PyArray_DESCR(op), ndim_output, new_dims, new_strides, PyArray_DATA(op), PyArray_ISWRITEABLE(op) ? NPY_ARRAY_WRITEABLE : 0, - (PyObject *)op, (PyObject *)op, - 0, 0); + (PyObject *)op, (PyObject *)op, 0); if (*ret == NULL) { return -1; @@ -532,10 +532,17 @@ unbuffered_loop_nop1_ndim2(NpyIter *iter) * Since the iterator wasn't tracking coordinates, the * loop provided by the iterator is in Fortran-order. */ - NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + int needs_api = NpyIter_IterationNeedsAPI(iter); + 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()){ + return -1; + } + ptr = ptrs[1][0] + strides[1][0]; ptrs[0][0] = ptrs[1][0] = ptr; ptr = ptrs[1][1] + strides[1][1]; @@ -586,11 +593,18 @@ unbuffered_loop_nop1_ndim3(NpyIter *iter) * Since the iterator wasn't tracking coordinates, the * loop provided by the iterator is in Fortran-order. */ - NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + int needs_api = NpyIter_IterationNeedsAPI(iter); + 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()){ + return -1; + } + ptr = ptrs[1][0] + strides[1][0]; ptrs[0][0] = ptrs[1][0] = ptr; ptr = ptrs[1][1] + strides[1][1]; @@ -643,10 +657,17 @@ unbuffered_loop_nop2_ndim2(NpyIter *iter) * Since the iterator wasn't tracking coordinates, the * loop provided by the iterator is in Fortran-order. */ - NPY_BEGIN_THREADS_THRESHOLDED(shape[1] * shape[0]); + int needs_api = NpyIter_IterationNeedsAPI(iter); + 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]); + if(needs_api && PyErr_Occurred()){ + return -1; + } + ptr = ptrs[1][0] + strides[1][0]; ptrs[0][0] = ptrs[1][0] = ptr; ptr = ptrs[1][1] + strides[1][1]; @@ -699,11 +720,18 @@ unbuffered_loop_nop2_ndim3(NpyIter *iter) * Since the iterator wasn't tracking coordinates, the * loop provided by the iterator is in Fortran-order. */ - NPY_BEGIN_THREADS_THRESHOLDED(shape[2] * shape[1] * shape[0]); + int needs_api = NpyIter_IterationNeedsAPI(iter); + 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]); + if(needs_api && PyErr_Occurred()){ + return -1; + } + ptr = ptrs[1][0] + strides[1][0]; ptrs[0][0] = ptrs[1][0] = ptr; ptr = ptrs[1][1] + strides[1][1]; @@ -1025,7 +1053,7 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop, goto fail; } - /* Initialize the output to all zeros */ + /* Initialize the output to all zeros or None*/ ret = NpyIter_GetOperandArray(iter)[nop]; if (PyArray_AssignZero(ret, NULL) < 0) { goto fail; |