Skip to content

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

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
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
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,8 @@ ignore = [
"FIX", # flake8-fixme
"ISC001", # Conflicts with formatter
"PYI041", # Use `float` instead of `int | float`
"TD002", # Missing author in TODO
"TD003", # Missing issue link for this TODO
]

[tool.ruff.lint.pydocstyle]
Expand Down
18 changes: 17 additions & 1 deletion src/array_api_typing/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,25 @@
"Array",
"HasArrayNamespace",
"HasDType",
"HasDevice",
"HasMatrixTranspose",
"HasNDim",
"HasShape",
"HasSize",
"HasTranspose",
"__version__",
"__version_tuple__",
)

from ._array import Array, HasArrayNamespace, HasDType
from ._array import (
Array,
HasArrayNamespace,
HasDevice,
HasDType,
HasMatrixTranspose,
HasNDim,
HasShape,
HasSize,
HasTranspose,
)
from ._version import version as __version__, version_tuple as __version_tuple__
132 changes: 130 additions & 2 deletions src/array_api_typing/_array.py
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)
Expand Down Expand Up @@ -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
"""Hardware device the array data resides on."""
...


class HasMatrixTranspose(Protocol):
"""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):
"""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, ...]:
"""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):
"""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):
"""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,
# ------- Methods ---------
HasArrayNamespace[NamespaceT_co],
# -------------------------
Protocol[DTypeT_co, NamespaceT_co],
):
Expand Down
32 changes: 30 additions & 2 deletions tests/integration/test_numpy1p0.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -49,5 +49,33 @@ a_ns: xpt.Array[Any, ModuleType] = nparr
# Note that `np.array_api` uses dtype objects, not dtype classes, so we can't
# type annotate specific dtypes like `np.float32` or `np.int32`.
_: xpt.Array[dtype[Any]] = nparr
_: xpt.Array[dtype[Any]] = nparr_i32
_: xpt.Array[dtype[Any]] = nparr_f32
x_f32: xpt.Array[dtype[Any]] = nparr_f32
x_i32: xpt.Array[dtype[Any]] = nparr_i32

# Check Attribute `.dtype`
_: dtype[Any] = x_f32.dtype
_: dtype[Any] = x_i32.dtype

# Check Attribute `.device`
_: object = x_f32.device
_: object = x_i32.device

# Check Attribute `.mT`
_: xpt.Array[dtype[Any]] = x_f32.mT
_: xpt.Array[dtype[Any]] = x_i32.mT

# Check Attribute `.ndim`
_: int = x_f32.ndim
_: int = x_i32.ndim

# Check Attribute `.shape`
_: tuple[int | None, ...] = x_f32.shape
_: tuple[int | None, ...] = x_i32.shape

# Check Attribute `.size`
_: int | None = x_f32.size
_: int | None = x_i32.size

# Check Attribute `.T`
_: xpt.Array[dtype[Any]] = x_f32.T
_: xpt.Array[dtype[Any]] = x_i32.T
46 changes: 41 additions & 5 deletions tests/integration/test_numpy2p0.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand All @@ -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`
Expand All @@ -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

# 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

# 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
Loading
Loading