Skip to content

Commit 2598ed9

Browse files
committed
✨: add HasArrayIMatmul protocol to support in-place matrix multiplication operator for array classes
1 parent 4645337 commit 2598ed9

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

src/array_api_typing/_array.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,22 @@ def __matmul__(self, other: Self, /) -> Self:
304304
...
305305

306306

307+
class HasArrayIMatmul(Protocol):
308+
"""Protocol for array classes that support the in-place matrix multiplication operator.""" # noqa: E501
309+
310+
def __imatmul__(self, other: Self, /) -> Self:
311+
"""Calculates the in-place matrix product for each element of an array instance with the respective element of the array other.
312+
313+
Args:
314+
other: multiplicand array. Must be compatible with self (see Broadcasting). Should have a numeric data type.
315+
316+
Returns:
317+
self, after performing the in-place matrix multiplication.
318+
319+
""" # noqa: E501
320+
...
321+
322+
307323
class HasArrayInvert(Protocol):
308324
"""Protocol for array classes that support the inverse operator."""
309325

@@ -513,6 +529,7 @@ class Array(
513529
HasArrayPow,
514530
HasArrayIPow,
515531
HasArrayMatmul,
532+
HasArrayIMatmul,
516533
HasArrayInvert,
517534
HasArrayAnd,
518535
HasArrayOr,

0 commit comments

Comments
 (0)