-
Notifications
You must be signed in to change notification settings - Fork 5
feat: HasX attributes #34
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
base: main
Are you sure you want to change the base?
Changes from all commits
4de00ae
591853b
b0b88cf
ed33843
e49a73a
0eba222
a5e3d5b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,10 +1,17 @@ | ||
__all__ = ( | ||
"Array", | ||
"HasArrayNamespace", | ||
"HasDType", | ||
"HasDevice", | ||
"HasMatrixTranspose", | ||
"HasNDim", | ||
"HasShape", | ||
"HasSize", | ||
"HasTranspose", | ||
) | ||
|
||
from types import ModuleType | ||
from typing import Literal, Protocol | ||
from typing import Literal, Protocol, Self | ||
from typing_extensions import TypeVar | ||
|
||
NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType) | ||
|
@@ -67,10 +74,131 @@ def dtype(self, /) -> DTypeT_co: | |
... | ||
|
||
|
||
class HasDevice(Protocol): | ||
"""Protocol for array classes that have a device attribute.""" | ||
|
||
@property | ||
def device(self) -> object: # TODO: more specific type | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I can see how this can be helpful for users if it would be generic. So I propose we either keep it private until that time, or to just make it generic right now. |
||
"""Hardware device the array data resides on.""" | ||
... | ||
|
||
|
||
class HasMatrixTranspose(Protocol): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd move this one closer to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. And the use-case of this is to require at least 2 dimensions, right? |
||
"""Protocol for array classes that have a matrix transpose attribute.""" | ||
|
||
@property | ||
def mT(self) -> Self: # noqa: N802 | ||
"""Transpose of a matrix (or a stack of matrices). | ||
|
||
If an array instance has fewer than two dimensions, an error should be | ||
raised. | ||
|
||
Returns: | ||
Self: array whose last two dimensions (axes) are permuted in reverse | ||
order relative to original array (i.e., for an array instance | ||
having shape `(..., M, N)`, the returned array must have shape | ||
`(..., N, M))`. The returned array must have the same data type | ||
as the original array. | ||
|
||
""" | ||
... | ||
|
||
|
||
class HasNDim(Protocol): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is there a situation where you'd wanna use
|
||
"""Protocol for array classes that have a number of dimensions attribute.""" | ||
|
||
@property | ||
def ndim(self) -> int: | ||
"""Number of array dimensions (axes). | ||
|
||
Returns: | ||
int: number of array dimensions (axes). | ||
|
||
""" | ||
... | ||
|
||
|
||
class HasShape(Protocol): | ||
"""Protocol for array classes that have a shape attribute.""" | ||
|
||
@property | ||
def shape(self) -> tuple[int | None, ...]: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think in a followup we should add a type parameter Array[..., UnknownDimsT=Never] Then we can define the alias types
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. what would There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
as in; what will it annotate There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The dim types, like def shape(self) -> tuple[int | UnknownDimsT, ...]: There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ah, ok, that makes sense. I was confused because it's a plural |
||
"""Shape of the array. | ||
|
||
Returns: | ||
tuple[int | None, ...]: array dimensions. An array dimension must be None | ||
if and only if a dimension is unknown. | ||
|
||
Notes: | ||
For array libraries having graph-based computational models, array | ||
dimensions may be unknown due to data-dependent operations (e.g., | ||
boolean indexing; `A[:, B > 0]`) and thus cannot be statically | ||
resolved without knowing array contents. | ||
|
||
""" | ||
... | ||
|
||
|
||
class HasSize(Protocol): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What's the use-case of this one? Is there anything this can help with, that |
||
"""Protocol for array classes that have a size attribute.""" | ||
|
||
@property | ||
def size(self) -> int | None: | ||
"""Number of elements in an array. | ||
|
||
Returns: | ||
int | None: number of elements in an array. The returned value must | ||
be `None` if and only if one or more array dimensions are | ||
unknown. | ||
|
||
Notes: | ||
This must equal the product of the array's dimensions. | ||
|
||
""" | ||
... | ||
|
||
|
||
class HasTranspose(Protocol): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ah nice; this could be used to annotate functions that exclusively accept 2d arrays, without having to deal with shape-typing :) |
||
"""Protocol for array classes that support the transpose operation.""" | ||
|
||
@property | ||
def T(self) -> Self: # noqa: N802 | ||
"""Transpose of the array. | ||
|
||
The array instance must be two-dimensional. If the array instance is not | ||
two-dimensional, an error should be raised. | ||
|
||
Returns: | ||
Self: two-dimensional array whose first and last dimensions (axes) | ||
are permuted in reverse order relative to original array. The | ||
returned array must have the same data type as the original | ||
array. | ||
|
||
Notes: | ||
Limiting the transpose to two-dimensional arrays (matrices) deviates | ||
from the NumPy et al practice of reversing all axes for arrays | ||
having more than two-dimensions. This is intentional, as reversing | ||
all axes was found to be problematic (e.g., conflicting with the | ||
mathematical definition of a transpose which is limited to matrices; | ||
not operating on batches of matrices; et cetera). In order to | ||
reverse all axes, one is recommended to use the functional | ||
`PermuteDims` interface found in this specification. | ||
|
||
""" | ||
... | ||
|
||
|
||
class Array( | ||
HasArrayNamespace[NamespaceT_co], | ||
# ------ Attributes ------- | ||
HasDType[DTypeT_co], | ||
HasDevice, | ||
HasMatrixTranspose, | ||
HasNDim, | ||
HasShape, | ||
HasSize, | ||
HasTranspose, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This |
||
# ------- Methods --------- | ||
HasArrayNamespace[NamespaceT_co], | ||
# ------------------------- | ||
Protocol[DTypeT_co, NamespaceT_co], | ||
): | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -11,12 +11,13 @@ import array_api_typing as xpt | |
# DType aliases | ||
F32: TypeAlias = np.float32 | ||
I32: TypeAlias = np.int32 | ||
B: TypeAlias = np.bool_ | ||
|
||
# Define NDArrays against which we can test the protocols | ||
nparr: npt.NDArray[Any] | ||
nparr_i32: npt.NDArray[I32] | ||
nparr_f32: npt.NDArray[F32] | ||
nparr_b: npt.NDArray[np.bool_] | ||
nparr_b: npt.NDArray[B] | ||
|
||
# ========================================================= | ||
# `xpt.HasArrayNamespace` | ||
|
@@ -42,7 +43,7 @@ _: xpt.HasArrayNamespace[dict[str, int]] = nparr # not caught | |
_: xpt.HasDType[Any] = nparr | ||
_: xpt.HasDType[np.dtype[I32]] = nparr_i32 | ||
_: xpt.HasDType[np.dtype[F32]] = nparr_f32 | ||
_: xpt.HasDType[np.dtype[np.bool_]] = nparr_b | ||
_: xpt.HasDType[np.dtype[B]] = nparr_b | ||
|
||
# ========================================================= | ||
# `xpt.Array` | ||
|
@@ -52,6 +53,41 @@ a_ns: xpt.Array[Any, ModuleType] = nparr | |
|
||
# Check DTypeT_co assignment | ||
_: xpt.Array[Any] = nparr | ||
_: xpt.Array[np.dtype[I32]] = nparr_i32 | ||
_: xpt.Array[np.dtype[F32]] = nparr_f32 | ||
_: xpt.Array[np.dtype[np.bool_]] = nparr_b | ||
x_f32: xpt.Array[np.dtype[F32]] = nparr_f32 | ||
x_i32: xpt.Array[np.dtype[I32]] = nparr_i32 | ||
x_b: xpt.Array[np.dtype[B]] = nparr_b | ||
|
||
# Check Attribute `.dtype` | ||
_: np.dtype[F32] = x_f32.dtype | ||
_: np.dtype[I32] = x_i32.dtype | ||
_: np.dtype[B] = x_b.dtype | ||
|
||
# Check Attribute `.device` | ||
_: object = x_f32.device | ||
_: object = x_i32.device | ||
_: object = x_b.device | ||
Comment on lines
+61
to
+68
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd go with |
||
|
||
# Check Attribute `.mT` | ||
_: xpt.Array[np.dtype[F32]] = x_f32.mT | ||
_: xpt.Array[np.dtype[I32]] = x_i32.mT | ||
_: xpt.Array[np.dtype[B]] = x_b.mT | ||
|
||
# Check Attribute `.ndim` | ||
_: int = x_f32.ndim | ||
_: int = x_i32.ndim | ||
_: int = x_b.ndim | ||
|
||
# Check Attribute `.shape` | ||
_: tuple[int | None, ...] = x_f32.shape | ||
_: tuple[int | None, ...] = x_i32.shape | ||
_: tuple[int | None, ...] = x_b.shape | ||
|
||
# Check Attribute `.size` | ||
_: int | None = x_f32.size | ||
_: int | None = x_i32.size | ||
_: int | None = x_b.size | ||
Comment on lines
+76
to
+88
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
||
# Check Attribute `.T` | ||
_: xpt.Array[np.dtype[F32]] = x_f32.T | ||
_: xpt.Array[np.dtype[I32]] = x_i32.T | ||
_: xpt.Array[np.dtype[B]] = x_b.T |
Uh oh!
There was an error while loading. Please reload this page.