Skip to content

Commit 522aaa1

Browse files
committed
✨ HasMatrixTranspose
1 parent 52645e9 commit 522aaa1

File tree

3 files changed

+32
-1
lines changed

3 files changed

+32
-1
lines changed

src/array_api_typing/_array.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
)
55

66
from types import ModuleType
7-
from typing import Literal, Protocol
7+
from typing import Literal, Protocol, Self
88
from typing_extensions import TypeVar
99

1010
NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType)
@@ -76,10 +76,32 @@ def device(self) -> object: # TODO: more specific type
7676
...
7777

7878

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

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)