Skip to content

Add implementation of dpnp.matrix_transpose() and .mT attribute for dpnp.array #2095

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 18 commits into from
Oct 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion doc/reference/linalg.rst
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,6 @@ Norms and other numbers
dpnp.trace
dpnp.linalg.trace (Array API compatible)


Solving linear equations
--------------------------

Expand Down
1 change: 1 addition & 0 deletions doc/reference/ndarray.rst
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ Other attributes
:nosignatures:

dpnp.ndarray.T
dpnp.ndarray.mT
dpnp.ndarray.real
dpnp.ndarray.imag
dpnp.ndarray.flat
Expand Down
43 changes: 43 additions & 0 deletions dpnp/dpnp_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,49 @@ def T(self):
"""View of the transposed array."""
return self.transpose()

@property
def mT(self):
"""
View of the matrix transposed array.

The matrix transpose is the transpose of the last two dimensions, even
if the array is of higher dimension.

Raises
------
ValueError
If the array is of dimension less than 2.

Examples
--------
>>> import dpnp as np
>>> a = np.array([[1, 2], [3, 4]])
>>> a
array([[1, 2],
[3, 4]])
>>> a.mT
array([[1, 3],
[2, 4]])

>>> a = np.arange(8).reshape((2, 2, 2))
>>> a
array([[[0, 1],
[2, 3]],
[[4, 5],
[6, 7]]])
>>> a.mT
array([[[0, 2],
[1, 3]],
[[4, 6],
[5, 7]]])

"""

if self.ndim < 2:
raise ValueError("matrix transpose with ndim < 2 is undefined")

return self._array_obj.mT

def to_device(self, target_device):
"""Transfer array to target device."""

Expand Down
53 changes: 53 additions & 0 deletions dpnp/dpnp_iface_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@
"flipud",
"hsplit",
"hstack",
"matrix_transpose",
"moveaxis",
"ndim",
"permute_dims",
Expand Down Expand Up @@ -1752,6 +1753,58 @@ def hstack(tup, *, dtype=None, casting="same_kind"):
return dpnp.concatenate(arrs, axis=1, dtype=dtype, casting=casting)


def matrix_transpose(x, /):
"""
Transposes a matrix (or a stack of matrices) `x`.

For full documentation refer to :obj:`numpy.matrix_transpose`.

Parameters
----------
x : (..., M, N) {dpnp.ndarray, usm_ndarray}
Input array with ``x.ndim >= 2`` and whose two innermost
dimensions form ``MxN`` matrices.

Returns
-------
out : dpnp.ndarray
An array containing the transpose for each matrix and having shape
(..., N, M).

See Also
--------
:obj:`dpnp.transpose` : Returns an array with axes transposed.
:obj:`dpnp.linalg.matrix_transpose` : Equivalent function.
:obj:`dpnp.ndarray.mT` : Equivalent method.

Examples
--------
>>> import dpnp as np
>>> a = np.array([[1, 2], [3, 4]])
>>> np.matrix_transpose(a)
array([[1, 3],
[2, 4]])

>>> b = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
>>> np.matrix_transpose(b)
array([[[1, 3],
[2, 4]],
[[5, 7],
[6, 8]]])

"""

usm_x = dpnp.get_usm_ndarray(x)
if usm_x.ndim < 2:
raise ValueError(
"Input array must be at least 2-dimensional, "
f"but it is {usm_x.ndim}"
)

usm_res = dpt.matrix_transpose(usm_x)
return dpnp_array._create_from_usm_ndarray(usm_res)


def moveaxis(a, source, destination):
"""
Move axes of an array to new positions. Other axes remain in their original
Expand Down
45 changes: 45 additions & 0 deletions dpnp/linalg/dpnp_iface_linalg.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@
"matrix_norm",
"matrix_power",
"matrix_rank",
"matrix_transpose",
"multi_dot",
"norm",
"outer",
Expand Down Expand Up @@ -1117,6 +1118,50 @@ def matrix_rank(A, tol=None, hermitian=False):
return dpnp_matrix_rank(A, tol=tol, hermitian=hermitian)


def matrix_transpose(x, /):
"""
Transposes a matrix (or a stack of matrices) `x`.

For full documentation refer to :obj:`numpy.linalg.matrix_transpose`.

Parameters
----------
x : (..., M, N) {dpnp.ndarray, usm_ndarray}
Input array with ``x.ndim >= 2`` and whose two innermost
dimensions form ``MxN`` matrices.

Returns
-------
out : dpnp.ndarray
An array containing the transpose for each matrix and having shape
(..., N, M).

See Also
--------
:obj:`dpnp.transpose` : Returns an array with axes transposed.
:obj:`dpnp.matrix_transpose` : Equivalent function.
:obj:`dpnp.ndarray.mT` : Equivalent method.

Examples
--------
>>> import dpnp as np
>>> a = np.array([[1, 2], [3, 4]])
>>> np.linalg.matrix_transpose(a)
array([[1, 3],
[2, 4]])

>>> b = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
>>> np.linalg.matrix_transpose(b)
array([[[1, 3],
[2, 4]],
[[5, 7],
[6, 8]]])

"""

return dpnp.matrix_transpose(x)


def multi_dot(arrays, *, out=None):
"""
Compute the dot product of two or more arrays in a single function call.
Expand Down
44 changes: 44 additions & 0 deletions tests/test_arraymanipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -556,6 +556,50 @@ def test_one_element(self):
assert_array_equal(res, a)


# numpy.matrix_transpose() is available since numpy >= 2.0
@testing.with_requires("numpy>=2.0")
class TestMatrixtranspose:
@pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
@pytest.mark.parametrize(
"shape",
[(3, 5), (4, 2), (2, 5, 2), (2, 3, 3, 6)],
ids=["(3,5)", "(4,2)", "(2,5,2)", "(2,3,3,6)"],
)
def test_matrix_transpose(self, dtype, shape):
a = numpy.arange(numpy.prod(shape), dtype=dtype).reshape(shape)
dp_a = dpnp.array(a)

expected = numpy.matrix_transpose(a)
result = dpnp.matrix_transpose(dp_a)

assert_allclose(result, expected)

@pytest.mark.parametrize(
"shape",
[(0, 0), (1, 0, 0), (0, 2, 2), (0, 1, 0, 4)],
ids=["(0,0)", "(1,0,0)", "(0,2,2)", "(0, 1, 0, 4)"],
)
def test_matrix_transpose_empty(self, shape):
a = numpy.empty(shape, dtype=dpnp.default_float_type())
dp_a = dpnp.array(a)

expected = numpy.matrix_transpose(a)
result = dpnp.matrix_transpose(dp_a)

assert_allclose(result, expected)

def test_matrix_transpose_errors(self):
a_dp = dpnp.array([[1, 2], [3, 4]], dtype="float32")

# unsupported type
a_np = dpnp.asnumpy(a_dp)
assert_raises(TypeError, dpnp.matrix_transpose, a_np)

# a.ndim < 2
a_dp_ndim_1 = a_dp.flatten()
assert_raises(ValueError, dpnp.matrix_transpose, a_dp_ndim_1)


class TestRollaxis:
data = [
(0, 0),
Expand Down
19 changes: 19 additions & 0 deletions tests/test_linalg.py
Original file line number Diff line number Diff line change
Expand Up @@ -2124,6 +2124,25 @@ def test_matrix_rank_errors(self):
)


# numpy.linalg.matrix_transpose() is available since numpy >= 2.0
@testing.with_requires("numpy>=2.0")
# dpnp.linalg.matrix_transpose() calls dpnp.matrix_transpose()
# 1 test to increase code coverage
def test_matrix_transpose():
a = numpy.arange(6).reshape((2, 3))
a_dp = inp.array(a)

expected = numpy.linalg.matrix_transpose(a)
result = inp.linalg.matrix_transpose(a_dp)

assert_allclose(expected, result)

with assert_raises_regex(
ValueError, "array must be at least 2-dimensional"
):
inp.linalg.matrix_transpose(a_dp[:, 0])


class TestNorm:
def setup_method(self):
numpy.random.seed(42)
Expand Down
33 changes: 33 additions & 0 deletions tests/test_ndarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from numpy.testing import assert_allclose, assert_array_equal

import dpnp
from tests.third_party.cupy import testing

from .helper import (
get_all_dtypes,
Expand Down Expand Up @@ -258,6 +259,38 @@ def test_array_as_index(shape, index_dtype):
assert a[tuple(ind_arr)] == a[1]


# numpy.ndarray.mT is available since numpy >= 2.0
@testing.with_requires("numpy>=2.0")
@pytest.mark.parametrize(
"shape",
[(3, 5), (2, 5, 2), (2, 3, 3, 6)],
ids=["(3,5)", "(2,5,2)", "(2,3,3,6)"],
)
def test_matrix_transpose(shape):
a = numpy.arange(numpy.prod(shape)).reshape(shape)
dp_a = dpnp.array(a)

expected = a.mT
result = dp_a.mT

assert_allclose(result, expected)

# result is a view of dp_a:
# changing result, modifies dp_a
first_elem = (0,) * dp_a.ndim

result[first_elem] = -1.0
assert dp_a[first_elem] == -1.0


@testing.with_requires("numpy>=2.0")
def test_matrix_transpose_error():
# 1D array
dp_a = dpnp.arange(6)
with pytest.raises(ValueError):
dp_a.mT


def test_ravel():
a = dpnp.ones((2, 2))
b = a.ravel()
Expand Down
Loading