Skip to content

Commit b0b88cf

Browse files
committed
✨ HasMatrixTranspose
1 parent 591853b commit b0b88cf

File tree

4 files changed

+35
-1
lines changed

4 files changed

+35
-1
lines changed

src/array_api_typing/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"HasArrayNamespace",
66
"HasDType",
77
"HasDevice",
8+
"HasMatrixTranspose",
89
"__version__",
910
"__version_tuple__",
1011
)
@@ -14,5 +15,6 @@
1415
HasArrayNamespace,
1516
HasDevice,
1617
HasDType,
18+
HasMatrixTranspose,
1719
)
1820
from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,11 @@
33
"HasArrayNamespace",
44
"HasDType",
55
"HasDevice",
6+
"HasMatrixTranspose",
67
)
78

89
from types import ModuleType
9-
from typing import Literal, Protocol
10+
from typing import Literal, Protocol, Self
1011
from typing_extensions import TypeVar
1112

1213
NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType)
@@ -78,10 +79,32 @@ def device(self) -> object: # TODO: more specific type
7879
...
7980

8081

82+
class HasMatrixTranspose(Protocol):
83+
"""Protocol for array classes that have a matrix transpose attribute."""
84+
85+
@property
86+
def mT(self) -> Self: # noqa: N802
87+
"""Transpose of a matrix (or a stack of matrices).
88+
89+
If an array instance has fewer than two dimensions, an error should be
90+
raised.
91+
92+
Returns:
93+
Self: array whose last two dimensions (axes) are permuted in reverse
94+
order relative to original array (i.e., for an array instance
95+
having shape `(..., M, N)`, the returned array must have shape
96+
`(..., N, M))`. The returned array must have the same data type
97+
as the original array.
98+
99+
"""
100+
...
101+
102+
81103
class Array(
82104
# ------ Attributes -------
83105
HasDType[DTypeT_co],
84106
HasDevice,
107+
HasMatrixTranspose,
85108
# ------- Methods ---------
86109
HasArrayNamespace[NamespaceT_co],
87110
# -------------------------

tests/integration/test_numpy1p0.pyi

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,3 +59,7 @@ _: dtype[Any] = x_i32.dtype
5959
# Check Attribute `.device`
6060
_: object = x_f32.device
6161
_: object = x_i32.device
62+
63+
# Check Attribute `.mT`
64+
_: xpt.Array[dtype[Any]] = x_f32.mT
65+
_: xpt.Array[dtype[Any]] = x_i32.mT

tests/integration/test_numpy2p0.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,8 @@ _: np.dtype[B] = x_b.dtype
6666
_: object = x_f32.device
6767
_: object = x_i32.device
6868
_: object = x_b.device
69+
70+
# Check Attribute `.mT`
71+
_: xpt.Array[np.dtype[F32]] = x_f32.mT
72+
_: xpt.Array[np.dtype[I32]] = x_i32.mT
73+
_: xpt.Array[np.dtype[B]] = x_b.mT

0 commit comments

Comments
 (0)