Merge branch 'master' into impl-searchsorted

Author: antonwolfy

dpnp/linalg/dpnp_iface_linalg.py

@@ -78,6 +78,8 @@
     "solve",
     "svd",
     "slogdet",
+    "tensorinv",
+    "tensorsolve",
 ]
 
 
@@ -930,3 +932,132 @@ def slogdet(a):
     check_stacked_square(a)
 
     return dpnp_slogdet(a)
+
+
+def tensorinv(a, ind=2):
+    """
+    Compute the 'inverse' of an N-dimensional array.
+
+    For full documentation refer to :obj:`numpy.linalg.tensorinv`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Tensor to `invert`. Its shape must be 'square', i. e.,
+        ``prod(a.shape[:ind]) == prod(a.shape[ind:])``.
+    ind : int, optional
+        Number of first indices that are involved in the inverse sum.
+        Must be a positive integer.
+        Default: 2.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The inverse of a tensor whose shape is equivalent to
+        ``a.shape[ind:] + a.shape[:ind]``.
+
+    See Also
+    --------
+    :obj:`dpnp.linalg.tensordot` : Compute tensor dot product along specified axes.
+    :obj:`dpnp.linalg.tensorsolve` : Solve the tensor equation ``a x = b`` for x.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.eye(4*6)
+    >>> a.shape = (4, 6, 8, 3)
+    >>> ainv = np.linalg.tensorinv(a, ind=2)
+    >>> ainv.shape
+    (8, 3, 4, 6)
+
+    >>> a = np.eye(4*6)
+    >>> a.shape = (24, 8, 3)
+    >>> ainv = np.linalg.tensorinv(a, ind=1)
+    >>> ainv.shape
+    (8, 3, 24)
+
+    """
+
+    dpnp.check_supported_arrays_type(a)
+
+    if ind <= 0:
+        raise ValueError("Invalid ind argument")
+
+    old_shape = a.shape
+    inv_shape = old_shape[ind:] + old_shape[:ind]
+    prod = numpy.prod(old_shape[ind:])
+    a = a.reshape(prod, -1)
+    a_inv = inv(a)
+
+    return a_inv.reshape(*inv_shape)
+
+
+def tensorsolve(a, b, axes=None):
+    """
+    Solve the tensor equation ``a x = b`` for x.
+
+    For full documentation refer to :obj:`numpy.linalg.tensorsolve`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
+        the shape of that sub-tensor of `a` consisting of the appropriate
+        number of its rightmost indices, and must be such that
+        ``prod(Q) == prod(b.shape)`` (in which sense `a` is said to be
+        'square').
+    b : {dpnp.ndarray, usm_ndarray}
+        Right-hand tensor, which can be of any shape.
+    axes : tuple of ints, optional
+        Axes in `a` to reorder to the right, before inversion.
+        If ``None`` , no reordering is done.
+        Default: ``None``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The tensor with shape ``Q`` such that ``b.shape + Q == a.shape``.
+
+    See Also
+    --------
+    :obj:`dpnp.linalg.tensordot` : Compute tensor dot product along specified axes.
+    :obj:`dpnp.linalg.tensorinv` : Compute the 'inverse' of an N-dimensional array.
+    :obj:`dpnp.einsum` : Evaluates the Einstein summation convention on the operands.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.eye(2*3*4)
+    >>> a.shape = (2*3, 4, 2, 3, 4)
+    >>> b = np.random.randn(2*3, 4)
+    >>> x = np.linalg.tensorsolve(a, b)
+    >>> x.shape
+    (2, 3, 4)
+    >>> np.allclose(np.tensordot(a, x, axes=3), b)
+    array([ True])
+
+    """
+
+    dpnp.check_supported_arrays_type(a, b)
+    a_ndim = a.ndim
+
+    if axes is not None:
+        all_axes = list(range(a_ndim))
+        for k in axes:
+            all_axes.remove(k)
+            all_axes.insert(a_ndim, k)
+        a = a.transpose(tuple(all_axes))
+
+    old_shape = a.shape[-(a_ndim - b.ndim) :]
+    prod = numpy.prod(old_shape)
+
+    if a.size != prod**2:
+        raise dpnp.linalg.LinAlgError(
+            "Input arrays must satisfy the requirement \
+            prod(a.shape[b.ndim:]) == prod(a.shape[:b.ndim])"
+        )
+
+    a = a.reshape(-1, prod)
+    b = b.ravel()
+    res = solve(a, b)
+    return res.reshape(old_shape)

tests/test_linalg.py

@@ -1459,3 +1459,86 @@ def test_pinv_errors(self):
         a_dp_q = inp.array(a_dp, sycl_queue=a_queue)
         rcond_dp_q = inp.array([0.5], dtype="float32", sycl_queue=rcond_queue)
         assert_raises(ValueError, inp.linalg.pinv, a_dp_q, rcond_dp_q)
+
+
+class TestTensorinv:
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    @pytest.mark.parametrize(
+        "shape, ind",
+        [
+            ((4, 6, 8, 3), 2),
+            ((24, 8, 3), 1),
+        ],
+        ids=[
+            "(4, 6, 8, 3)",
+            "(24, 8, 3)",
+        ],
+    )
+    def test_tensorinv(self, dtype, shape, ind):
+        a = numpy.eye(24, dtype=dtype).reshape(shape)
+        a_dp = inp.array(a)
+
+        ainv = numpy.linalg.tensorinv(a, ind=ind)
+        ainv_dp = inp.linalg.tensorinv(a_dp, ind=ind)
+
+        assert ainv.shape == ainv_dp.shape
+        assert_dtype_allclose(ainv_dp, ainv)
+
+    def test_test_tensorinv_errors(self):
+        a_dp = inp.eye(24, dtype="float32").reshape(4, 6, 8, 3)
+
+        # unsupported type `a`
+        a_np = inp.asnumpy(a_dp)
+        assert_raises(TypeError, inp.linalg.pinv, a_np)
+
+        # unsupported type `ind`
+        assert_raises(TypeError, inp.linalg.tensorinv, a_dp, 2.0)
+        assert_raises(TypeError, inp.linalg.tensorinv, a_dp, [2.0])
+        assert_raises(ValueError, inp.linalg.tensorinv, a_dp, -1)
+
+        # non-square
+        assert_raises(inp.linalg.LinAlgError, inp.linalg.tensorinv, a_dp, 1)
+
+
+class TestTensorsolve:
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    @pytest.mark.parametrize(
+        "axes",
+        [None, (1,), (2,)],
+        ids=[
+            "None",
+            "(1,)",
+            "(2,)",
+        ],
+    )
+    def test_tensorsolve_axes(self, dtype, axes):
+        a = numpy.eye(12).reshape(12, 3, 4).astype(dtype)
+        b = numpy.ones(a.shape[0], dtype=dtype)
+
+        a_dp = inp.array(a)
+        b_dp = inp.array(b)
+
+        res_np = numpy.linalg.tensorsolve(a, b, axes=axes)
+        res_dp = inp.linalg.tensorsolve(a_dp, b_dp, axes=axes)
+
+        assert res_np.shape == res_dp.shape
+        assert_dtype_allclose(res_dp, res_np)
+
+    def test_tensorsolve_errors(self):
+        a_dp = inp.eye(24, dtype="float32").reshape(4, 6, 8, 3)
+        b_dp = inp.ones(a_dp.shape[:2], dtype="float32")
+
+        # unsupported type `a` and `b`
+        a_np = inp.asnumpy(a_dp)
+        b_np = inp.asnumpy(b_dp)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_np, b_dp)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, b_np)
+
+        # unsupported type `axes`
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, 2.0)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, -2)
+
+        # incorrect axes
+        assert_raises(
+            inp.linalg.LinAlgError, inp.linalg.tensorsolve, a_dp, b_dp, (1,)
+        )

tests/test_sycl_queue.py

@@ -1896,3 +1896,42 @@ def test_pinv(shape, hermitian, rcond_as_array, device):
     B_queue = B_result.sycl_queue
 
     assert_sycl_queue_equal(B_queue, a_dp.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_tensorinv(device):
+    a_np = numpy.eye(12).reshape(12, 4, 3)
+    a_dp = dpnp.array(a_np, device=device)
+
+    result = dpnp.linalg.tensorinv(a_dp, ind=1)
+    expected = numpy.linalg.tensorinv(a_np, ind=1)
+    assert_dtype_allclose(result, expected)
+
+    result_queue = result.sycl_queue
+
+    assert_sycl_queue_equal(result_queue, a_dp.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_tensorsolve(device):
+    a_np = numpy.random.randn(3, 2, 6).astype(dpnp.default_float_type())
+    b_np = numpy.ones(a_np.shape[:2], dtype=a_np.dtype)
+
+    a_dp = dpnp.array(a_np, device=device)
+    b_dp = dpnp.array(b_np, device=device)
+
+    result = dpnp.linalg.tensorsolve(a_dp, b_dp)
+    expected = numpy.linalg.tensorsolve(a_np, b_np)
+    assert_dtype_allclose(result, expected)
+
+    result_queue = result.sycl_queue
+
+    assert_sycl_queue_equal(result_queue, a_dp.sycl_queue)

tests/test_usm_type.py

@@ -1041,3 +1041,25 @@ def test_qr(shape, mode, usm_type):
 
         assert a.usm_type == dp_q.usm_type
         assert a.usm_type == dp_r.usm_type
+
+
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+def test_tensorinv(usm_type):
+    a = dp.eye(12, usm_type=usm_type).reshape(12, 4, 3)
+    ainv = dp.linalg.tensorinv(a, ind=1)
+
+    assert a.usm_type == ainv.usm_type
+
+
+@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_b", list_of_usm_types, ids=list_of_usm_types)
+def test_tensorsolve(usm_type_a, usm_type_b):
+    data = numpy.random.randn(3, 2, 6)
+    a = dp.array(data, usm_type=usm_type_a)
+    b = dp.ones(a.shape[:2], dtype=a.dtype, usm_type=usm_type_b)
+
+    result = dp.linalg.tensorsolve(a, b)
+
+    assert a.usm_type == usm_type_a
+    assert b.usm_type == usm_type_b
+    assert result.usm_type == du.get_coerced_usm_type([usm_type_a, usm_type_b])

tests/third_party/cupy/linalg_tests/test_solve.py

@@ -95,6 +95,26 @@ def test_invalid_shape(self):
         self.check_shape((0, 3, 4), (3,), linalg_errors)
 
 
+@testing.parameterize(
+    *testing.product(
+        {
+            "a_shape": [(2, 3, 6), (3, 4, 4, 3)],
+            "axes": [None, (0, 2)],
+        }
+    )
+)
+@testing.fix_random()
+class TestTensorSolve(unittest.TestCase):
+    @testing.for_dtypes("ifdFD")
+    @testing.numpy_cupy_allclose(atol=0.02, type_check=has_support_aspect64())
+    def test_tensorsolve(self, xp, dtype):
+        a_shape = self.a_shape
+        b_shape = self.a_shape[:2]
+        a = testing.shaped_random(a_shape, xp, dtype=dtype, seed=0)
+        b = testing.shaped_random(b_shape, xp, dtype=dtype, seed=1)
+        return xp.linalg.tensorsolve(a, b, axes=self.axes)
+
+
 @testing.parameterize(
     *testing.product(
         {
@@ -208,3 +228,47 @@ def test_pinv_size_0(self):
         self.check_x((0, 0), rcond=1e-15)
         self.check_x((0, 2, 3), rcond=1e-15)
         self.check_x((2, 0, 3), rcond=1e-15)
+
+
+class TestTensorInv(unittest.TestCase):
+    @testing.for_dtypes("ifdFD")
+    @_condition.retry(10)
+    def check_x(self, a_shape, ind, dtype):
+        a_cpu = numpy.random.randint(0, 10, size=a_shape).astype(dtype)
+        a_gpu = cupy.asarray(a_cpu)
+        a_gpu_copy = a_gpu.copy()
+        result_cpu = numpy.linalg.tensorinv(a_cpu, ind=ind)
+        result_gpu = cupy.linalg.tensorinv(a_gpu, ind=ind)
+        assert_dtype_allclose(result_gpu, result_cpu)
+        testing.assert_array_equal(a_gpu_copy, a_gpu)
+
+    def check_shape(self, a_shape, ind):
+        a = cupy.random.rand(*a_shape)
+        with self.assertRaises(
+            (numpy.linalg.LinAlgError, cupy.linalg.LinAlgError)
+        ):
+            cupy.linalg.tensorinv(a, ind=ind)
+
+    def check_ind(self, a_shape, ind):
+        a = cupy.random.rand(*a_shape)
+        with self.assertRaises(ValueError):
+            cupy.linalg.tensorinv(a, ind=ind)
+
+    def test_tensorinv(self):
+        self.check_x((12, 3, 4), ind=1)
+        self.check_x((3, 8, 24), ind=2)
+        self.check_x((18, 3, 3, 2), ind=1)
+        self.check_x((1, 4, 2, 2), ind=2)
+        self.check_x((2, 3, 5, 30), ind=3)
+        self.check_x((24, 2, 2, 3, 2), ind=1)
+        self.check_x((3, 4, 2, 3, 2), ind=2)
+        self.check_x((1, 2, 3, 2, 3), ind=3)
+        self.check_x((3, 2, 1, 2, 12), ind=4)
+
+    def test_invalid_shape(self):
+        self.check_shape((2, 3, 4), ind=1)
+        self.check_shape((1, 2, 3, 4), ind=3)
+
+    def test_invalid_index(self):
+        self.check_ind((12, 3, 4), ind=-1)
+        self.check_ind((18, 3, 3, 2), ind=0)