Update dpnp.linalg.matrix_power() implementation

dpnp/linalg/dpnp_iface_linalg.py

@@ -51,6 +51,7 @@
     dpnp_det,
     dpnp_eigh,
     dpnp_inv,
+    dpnp_matrix_power,
     dpnp_matrix_rank,
     dpnp_multi_dot,
     dpnp_pinv,
@@ -370,33 +371,65 @@ def inv(a):
     return dpnp_inv(a)
 
 
-def matrix_power(input, count):
+def matrix_power(a, n):
     """
-    Raise a square matrix to the (integer) power `count`.
+    Raise a square matrix to the (integer) power `n`.
+
+    For full documentation refer to :obj:`numpy.linalg.matrix_power`.
 
     Parameters
     ----------
-    input : sequence of array_like
+    a : (..., M, M) {dpnp.ndarray, usm_ndarray}
+        Matrix to be "powered".
+    n : int
+        The exponent can be any integer or long integer, positive, negative, or zero.
 
     Returns
     -------
-    output : array
-        Returns the dot product of the supplied arrays.
+    a**n : (..., M, M) dpnp.ndarray
+        The return value is the same shape and type as `M`;
+        if the exponent is positive or zero then the type of the
+        elements is the same as those of `M`. If the exponent is
+        negative the elements are floating-point.
 
-    See Also
-    --------
-    :obj:`numpy.linalg.matrix_power`
+    >>> import dpnp as np
+    >>> i = np.array([[0, 1], [-1, 0]]) # matrix equiv. of the imaginary unit
+    >>> np.linalg.matrix_power(i, 3) # should = -i
+    array([[ 0, -1],
+           [ 1,  0]])
+    >>> np.linalg.matrix_power(i, 0)
+    array([[1, 0],
+           [0, 1]])
+    >>> np.linalg.matrix_power(i, -3) # should = 1/(-i) = i, but w/ f.p. elements
+    array([[ 0.,  1.],
+           [-1.,  0.]])
+
+    Somewhat more sophisticated example
+
+    >>> q = np.zeros((4, 4))
+    >>> q[0:2, 0:2] = -i
+    >>> q[2:4, 2:4] = i
+    >>> q # one of the three quaternion units not equal to 1
+    array([[ 0., -1.,  0.,  0.],
+           [ 1.,  0.,  0.,  0.],
+           [ 0.,  0.,  0.,  1.],
+           [ 0.,  0., -1.,  0.]])
+    >>> np.linalg.matrix_power(q, 2) # = -np.eye(4)
+    array([[-1.,  0.,  0.,  0.],
+           [ 0., -1.,  0.,  0.],
+           [ 0.,  0., -1.,  0.],
+           [ 0.,  0.,  0., -1.]])
 
     """
 
-    if not use_origin_backend() and count > 0:
-        result = input
-        for _ in range(count - 1):
-            result = dpnp.matmul(result, input)
+    dpnp.check_supported_arrays_type(a)
+    check_stacked_2d(a)
+    check_stacked_square(a)
 
-        return result
+    if not isinstance(n, int):
+        raise TypeError("exponent must be an integer")
 
-    return call_origin(numpy.linalg.matrix_power, input, count)
+    return dpnp_matrix_power(a, n)
 
 
 def matrix_rank(A, tol=None, hermitian=False):

dpnp/linalg/dpnp_utils_linalg.py

@@ -40,6 +40,7 @@
     "dpnp_det",
     "dpnp_eigh",
     "dpnp_inv",
+    "dpnp_matrix_power",
     "dpnp_matrix_rank",
     "dpnp_multi_dot",
     "dpnp_pinv",
@@ -532,6 +533,46 @@ def _stacked_identity(
     return x
 
 
+def _stacked_identity_like(x):
+    """
+    Create stacked identity matrices based on the shape and properties of `x`.
+
+    Parameters
+    ----------
+    x : dpnp.ndarray
+        Input array based on whose properties (shape, data type, USM type and SYCL queue)
+        the identity matrices will be created.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        Array of stacked `n x n` identity matrices,
+        where `n` is the size of the last dimension of `x`.
+        The returned array has the same shape, data type, USM type
+        and uses the same SYCL queue as `x`, if applicable.
+
+    Example
+    -------
+    >>> import dpnp
+    >>> x = dpnp.zeros((2, 3, 3), dtype=dpnp.int64)
+    >>> _stacked_identity_like(x)
+    array([[[1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1]],
+
+           [[1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1]]], dtype=int32)
+
+    """
+
+    n = x.shape[-1]
+    idx = dpnp.arange(n, usm_type=x.usm_type, sycl_queue=x.sycl_queue)
+    x = dpnp.zeros_like(x)
+    x[..., idx, idx] = 1
+    return x
+
+
 def _triu_inplace(a, host_tasks, depends=None):
     """
     _triu_inplace(a, host_tasks, depends=None)
@@ -1082,6 +1123,41 @@ def dpnp_inv(a):
     return b_f
 
 
+def dpnp_matrix_power(a, n):
+    """
+    dpnp_matrix_power(a, n)
+
+    Raise a square matrix to the (integer) power `n`.
+
+    """
+
+    if n == 0:
+        return _stacked_identity_like(a)
+    elif n < 0:
+        a = dpnp.linalg.inv(a)
+        n *= -1
+
+    if n <= 3:
+        if n == 1:
+            return a
+        elif n == 2:
+            return dpnp.matmul(a, a)
+        else:
+            return dpnp.matmul(dpnp.matmul(a, a), a)
+
+    # binary decomposition to reduce the number of matrix
+    # multiplications for n > 3.
+    # `result` will hold the final matrix power,
+    # while `acc` serves as an accumulator for the intermediate matrix powers.
+    result, acc = None, None
+    for b in numpy.base_repr(n)[::-1]:
+        acc = a if acc is None else dpnp.matmul(acc, acc)
+        if b == "1":
+            result = acc if result is None else dpnp.matmul(result, acc)
+
+    return result
+
+
 def dpnp_matrix_rank(A, tol=None, hermitian=False):
     """
     dpnp_matrix_rank(A, tol=None, hermitian=False)

tests/skipped_tests.tbl

@@ -332,10 +332,6 @@ tests/third_party/cupy/linalg_tests/test_einsum.py::TestListArgEinSumError::test
 tests/third_party/cupy/linalg_tests/test_einsum.py::TestListArgEinSumError::test_invalid_sub1
 tests/third_party/cupy/linalg_tests/test_einsum.py::TestListArgEinSumError::test_too_many_dims3
 
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_invlarge
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_large
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_of_two
-
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_broadcast_not_allowed
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_diff_dtypes_is_equal
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_diff_dtypes_not_equal

tests/skipped_tests_gpu.tbl

@@ -434,10 +434,6 @@ tests/third_party/cupy/linalg_tests/test_einsum.py::TestEinSumUnaryOperationWith
 tests/third_party/cupy/linalg_tests/test_einsum.py::TestEinSumUnaryOperationWithScalar::test_scalar_int
 tests/third_party/cupy/linalg_tests/test_einsum.py::TestListArgEinSumError::test_invalid_sub1
 
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_invlarge
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_large
-tests/third_party/cupy/linalg_tests/test_product.py::TestMatrixPower::test_matrix_power_of_two
-
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_broadcast_not_allowed
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_diff_dtypes_is_equal
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_diff_dtypes_not_equal

tests/test_linalg.py

@@ -566,6 +566,54 @@ def test_inv_errors(self):
         assert_raises(inp.linalg.LinAlgError, inp.linalg.inv, a_dp)
 
 
+class TestMatrixPower:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
+    @pytest.mark.parametrize(
+        "data, power",
+        [
+            (
+                numpy.block(
+                    [
+                        [numpy.eye(2), numpy.zeros((2, 2))],
+                        [numpy.zeros((2, 2)), numpy.eye(2) * 2],
+                    ]
+                ),
+                3,
+            ),  # Block-diagonal matrix
+            (numpy.eye(3, k=1) + numpy.eye(3), 3),  # Non-diagonal matrix
+            (
+                numpy.eye(3, k=1) + numpy.eye(3),
+                -3,
+            ),  # Inverse of non-diagonal matrix
+        ],
+    )
+    def test_matrix_power(self, data, power, dtype):
+        a = data.astype(dtype)
+        a_dp = inp.array(a)
+
+        result = inp.linalg.matrix_power(a_dp, power)
+        expected = numpy.linalg.matrix_power(a, power)
+
+        assert_dtype_allclose(result, expected)
+
+    def test_matrix_power_errors(self):
+        a_dp = inp.eye(4, dtype="float32")
+
+        # unsupported type `a`
+        a_np = inp.asnumpy(a_dp)
+        assert_raises(TypeError, inp.linalg.matrix_power, a_np, 2)
+
+        # unsupported type `power`
+        assert_raises(TypeError, inp.linalg.matrix_power, a_dp, 1.5)
+        assert_raises(TypeError, inp.linalg.matrix_power, a_dp, [2])
+
+        # not invertible
+        noninv = inp.array([[1, 0], [0, 0]])
+        assert_raises(
+            inp.linalg.LinAlgError, inp.linalg.matrix_power, noninv, -1
+        )
+
+
 class TestMatrixRank:
     @pytest.mark.parametrize("dtype", get_all_dtypes())
     @pytest.mark.parametrize(

tests/third_party/cupy/linalg_tests/test_product.py

@@ -430,7 +430,6 @@ def test_tensordot_zero_length(self, xp, dtype):
 
 
 class TestMatrixPower(unittest.TestCase):
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()
     def test_matrix_power_0(self, xp, dtype):
@@ -455,23 +454,20 @@ def test_matrix_power_3(self, xp, dtype):
         a = testing.shaped_arange((3, 3), xp, dtype)
         return xp.linalg.matrix_power(a, 3)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_float_dtypes(no_float16=True)
     @testing.numpy_cupy_allclose(rtol=1e-5)
     def test_matrix_power_inv1(self, xp, dtype):
         a = testing.shaped_arange((3, 3), xp, dtype)
         a = a * a % 30
         return xp.linalg.matrix_power(a, -1)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_float_dtypes(no_float16=True)
     @testing.numpy_cupy_allclose(rtol=1e-5)
     def test_matrix_power_inv2(self, xp, dtype):
         a = testing.shaped_arange((3, 3), xp, dtype)
         a = a * a % 30
         return xp.linalg.matrix_power(a, -2)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_float_dtypes(no_float16=True)
     @testing.numpy_cupy_allclose(rtol=1e-4)
     def test_matrix_power_inv3(self, xp, dtype):