summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/einsumfunc.py4
-rw-r--r--numpy/core/shape_base.py6
-rw-r--r--numpy/core/src/multiarray/iterators.c108
-rw-r--r--numpy/core/src/multiarray/iterators.h11
-rw-r--r--numpy/linalg/tests/test_linalg.py41
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))