diff options
Diffstat (limited to 'numpy/linalg/linalg.py')
| -rw-r--r-- | numpy/linalg/linalg.py | 20 |
1 files changed, 16 insertions, 4 deletions
diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py index 5f04620ac..ccc437663 100644 --- a/numpy/linalg/linalg.py +++ b/numpy/linalg/linalg.py @@ -537,6 +537,8 @@ def matrix_power(a, n): of the same shape as M is returned. If ``n < 0``, the inverse is computed and then raised to the ``abs(n)``. + .. note:: Stacks of object matrices are not currently supported. + Parameters ---------- a : (..., M, M) array_like @@ -599,6 +601,16 @@ def matrix_power(a, n): except TypeError: raise TypeError("exponent must be an integer") + # Fall back on dot for object arrays. Object arrays are not supported by + # the current implementation of matmul using einsum + if a.dtype != object: + fmatmul = matmul + elif a.ndim == 2: + fmatmul = dot + else: + raise NotImplementedError( + "matrix_power not supported for stacks of object arrays") + if n == 0: a = empty_like(a) a[...] = eye(a.shape[-2], dtype=a.dtype) @@ -613,20 +625,20 @@ def matrix_power(a, n): return a elif n == 2: - return matmul(a, a) + return fmatmul(a, a) elif n == 3: - return matmul(matmul(a, a), a) + return fmatmul(fmatmul(a, a), a) # Use binary decomposition to reduce the number of matrix multiplications. # Here, we iterate over the bits of n, from LSB to MSB, raise `a` to # increasing powers of 2, and multiply into the result as needed. z = result = None while n > 0: - z = a if z is None else matmul(z, z) + z = a if z is None else fmatmul(z, z) n, bit = divmod(n, 2) if bit: - result = z if result is None else matmul(result, z) + result = z if result is None else fmatmul(result, z) return result |