diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/einsumfunc.py | 4 | ||||
| -rw-r--r-- | numpy/core/shape_base.py | 6 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/iterators.c | 108 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/iterators.h | 11 | ||||
| -rw-r--r-- | numpy/linalg/tests/test_linalg.py | 41 |
5 files changed, 48 insertions, 122 deletions
diff --git a/numpy/core/einsumfunc.py b/numpy/core/einsumfunc.py index 3ffb152e1..d9f88cb1c 100644 --- a/numpy/core/einsumfunc.py +++ b/numpy/core/einsumfunc.py @@ -700,7 +700,7 @@ def _einsum_path_dispatcher(*operands, **kwargs): return operands -@array_function_dispatch(_einsum_path_dispatcher) +@array_function_dispatch(_einsum_path_dispatcher, module='numpy') def einsum_path(*operands, **kwargs): """ einsum_path(subscripts, *operands, optimize='greedy') @@ -1001,7 +1001,7 @@ def _einsum_dispatcher(*operands, **kwargs): # Rewrite einsum to handle different cases -@array_function_dispatch(_einsum_dispatcher) +@array_function_dispatch(_einsum_dispatcher, module='numpy') def einsum(*operands, **kwargs): """ einsum(subscripts, *operands, out=None, dtype=None, order='K', diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py index c9f8ebccb..71a23f438 100644 --- a/numpy/core/shape_base.py +++ b/numpy/core/shape_base.py @@ -7,9 +7,13 @@ import functools import operator from . import numeric as _nx +from . import overrides from .numeric import array, asanyarray, newaxis from .multiarray import normalize_axis_index -from .overrides import array_function_dispatch + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') def _atleast_1d_dispatcher(*arys): diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c index 3e3248f53..a3bc8e742 100644 --- a/numpy/core/src/multiarray/iterators.c +++ b/numpy/core/src/multiarray/iterators.c @@ -92,114 +92,6 @@ parse_index_entry(PyObject *op, npy_intp *step_size, } -/* - * Parses an index that has no fancy indexing. Populates - * out_dimensions, out_strides, and out_offset. - */ -NPY_NO_EXPORT int -parse_index(PyArrayObject *self, PyObject *op, - npy_intp *out_dimensions, - npy_intp *out_strides, - npy_intp *out_offset, - int check_index) -{ - int i, j, n; - int nd_old, nd_new, n_add, n_ellipsis; - npy_intp n_steps, start, offset, step_size; - PyObject *op1 = NULL; - int is_slice; - - if (PySlice_Check(op) || op == Py_Ellipsis || op == Py_None) { - n = 1; - op1 = op; - Py_INCREF(op); - /* this relies on the fact that n==1 for loop below */ - is_slice = 1; - } - else { - if (!PySequence_Check(op)) { - PyErr_SetString(PyExc_IndexError, - "index must be either an int " - "or a sequence"); - return -1; - } - n = PySequence_Length(op); - is_slice = 0; - } - - nd_old = nd_new = 0; - - offset = 0; - for (i = 0; i < n; i++) { - if (!is_slice) { - op1 = PySequence_GetItem(op, i); - if (op1 == NULL) { - return -1; - } - } - start = parse_index_entry(op1, &step_size, &n_steps, - nd_old < PyArray_NDIM(self) ? - PyArray_DIMS(self)[nd_old] : 0, - nd_old, check_index ? - nd_old < PyArray_NDIM(self) : 0); - Py_DECREF(op1); - if (start == -1) { - break; - } - if (n_steps == NEWAXIS_INDEX) { - out_dimensions[nd_new] = 1; - out_strides[nd_new] = 0; - nd_new++; - } - else if (n_steps == ELLIPSIS_INDEX) { - for (j = i + 1, n_ellipsis = 0; j < n; j++) { - op1 = PySequence_GetItem(op, j); - if (op1 == Py_None) { - n_ellipsis++; - } - Py_DECREF(op1); - } - n_add = PyArray_NDIM(self)-(n-i-n_ellipsis-1+nd_old); - if (n_add < 0) { - PyErr_SetString(PyExc_IndexError, "too many indices"); - return -1; - } - for (j = 0; j < n_add; j++) { - out_dimensions[nd_new] = PyArray_DIMS(self)[nd_old]; - out_strides[nd_new] = PyArray_STRIDES(self)[nd_old]; - nd_new++; nd_old++; - } - } - else { - if (nd_old >= PyArray_NDIM(self)) { - PyErr_SetString(PyExc_IndexError, "too many indices"); - return -1; - } - offset += PyArray_STRIDES(self)[nd_old]*start; - nd_old++; - if (n_steps != SINGLE_INDEX) { - out_dimensions[nd_new] = n_steps; - out_strides[nd_new] = step_size * - PyArray_STRIDES(self)[nd_old-1]; - nd_new++; - } - } - } - if (i < n) { - return -1; - } - n_add = PyArray_NDIM(self)-nd_old; - for (j = 0; j < n_add; j++) { - out_dimensions[nd_new] = PyArray_DIMS(self)[nd_old]; - out_strides[nd_new] = PyArray_STRIDES(self)[nd_old]; - nd_new++; - nd_old++; - } - *out_offset = offset; - return nd_new; -} - - /*********************** Element-wise Array Iterator ***********************/ /* Aided by Peter J. Verveer's nd_image package and numpy's arraymap ****/ /* and Python's array iterator ***/ diff --git a/numpy/core/src/multiarray/iterators.h b/numpy/core/src/multiarray/iterators.h index 04f57c885..376dc154a 100644 --- a/numpy/core/src/multiarray/iterators.h +++ b/numpy/core/src/multiarray/iterators.h @@ -1,17 +1,6 @@ #ifndef _NPY_ARRAYITERATORS_H_ #define _NPY_ARRAYITERATORS_H_ -/* - * Parses an index that has no fancy indexing. Populates - * out_dimensions, out_strides, and out_offset. - */ -NPY_NO_EXPORT int -parse_index(PyArrayObject *self, PyObject *op, - npy_intp *out_dimensions, - npy_intp *out_strides, - npy_intp *out_offset, - int check_index); - NPY_NO_EXPORT PyObject *iter_subscript(PyArrayIterObject *, PyObject *); diff --git a/numpy/linalg/tests/test_linalg.py b/numpy/linalg/tests/test_linalg.py index 0e94c2633..836681039 100644 --- a/numpy/linalg/tests/test_linalg.py +++ b/numpy/linalg/tests/test_linalg.py @@ -1915,3 +1915,44 @@ class TestMultiDot(object): def test_too_few_input_arrays(self): assert_raises(ValueError, multi_dot, []) assert_raises(ValueError, multi_dot, [np.random.random((3, 3))]) + + +class TestTensorinv(object): + + @pytest.mark.parametrize("arr, ind", [ + (np.ones((4, 6, 8, 2)), 2), + (np.ones((3, 3, 2)), 1), + ]) + def test_non_square_handling(self, arr, ind): + with assert_raises(LinAlgError): + linalg.tensorinv(arr, ind=ind) + + @pytest.mark.parametrize("shape, ind", [ + # examples from docstring + ((4, 6, 8, 3), 2), + ((24, 8, 3), 1), + ]) + def test_tensorinv_shape(self, shape, ind): + a = np.eye(24) + a.shape = shape + ainv = linalg.tensorinv(a=a, ind=ind) + expected = a.shape[ind:] + a.shape[:ind] + actual = ainv.shape + assert_equal(actual, expected) + + @pytest.mark.parametrize("ind", [ + 0, -2, + ]) + def test_tensorinv_ind_limit(self, ind): + a = np.eye(24) + a.shape = (4, 6, 8, 3) + with assert_raises(ValueError): + linalg.tensorinv(a=a, ind=ind) + + def test_tensorinv_result(self): + # mimic a docstring example + a = np.eye(24) + a.shape = (24, 8, 3) + ainv = linalg.tensorinv(a, ind=1) + b = np.ones(24) + assert_allclose(np.tensordot(ainv, b, 1), np.linalg.tensorsolve(a, b)) |
