Add Array API inspection utilities (#689)

kgryte · web-flow · commit cc11aa344d8d · 2024-02-14T18:37:07.000-05:00
diff --git a/spec/draft/API_specification/index.rst b/spec/draft/API_specification/index.rst
@@ -29,6 +29,7 @@ A conforming implementation of the array API standard must provide and support t
    function_and_method_signatures
    indexing
    indexing_functions
+   inspection
    linear_algebra_functions
    manipulation_functions
    searching_functions
diff --git a/spec/draft/API_specification/inspection.rst b/spec/draft/API_specification/inspection.rst
@@ -0,0 +1,40 @@
+Inspection
+==========
+
+    Array API specification for namespace inspection utilities.
+
+A conforming implementation of the array API standard must provide and support the following functions and associated inspection APIs.
+
+
+Objects in API
+--------------
+
+.. currentmodule:: array_api.info
+
+..
+  NOTE: please keep the functions in alphabetical order
+
+.. autosummary::
+   :toctree: generated
+   :template: method.rst
+
+   __array_namespace_info__
+
+
+Inspection APIs
+---------------
+
+In the namespace (or class) returned by ``__array_namespace_info__``, a conforming implementation of the array API standard must provide and support the following functions (or methods) for programmatically querying data type and device support, capabilities, and other specification-defined implementation-specific behavior, as documented in the functions described below.
+
+..
+  NOTE: please keep the functions in alphabetical order
+
+.. autosummary::
+   :toctree: generated
+   :template: method.rst
+
+   capabilities
+   default_device
+   default_dtypes
+   devices
+   dtypes
diff --git a/spec/draft/design_topics/device_support.rst b/spec/draft/design_topics/device_support.rst
@@ -39,7 +39,7 @@ into library code which may have to do the following:
 Syntax for device assignment
 ----------------------------
 
-The array API will offer the following syntax for device assignment and
+The array API provides the following syntax for device assignment and
 cross-device data transfer:
 
 1. A ``.device`` property on the array object, which returns a ``Device`` object
@@ -52,19 +52,20 @@ cross-device data transfer:
 3. A ``.to_device`` method on the array object to copy an array to a different device.
 
 .. note::
-   In the current API standard, the only way to obtain a ``Device`` object is from the
-   ``.device`` property on the array object. The standard does **not** include a universal
-   ``Device`` object recognized by all compliant libraries. Accordingly, the standard does
-   not provide a means of instantiating a ``Device`` object to point to a specific physical or
-   logical device.
-
-   The choice to not include a standardized ``Device`` object may be revisited in a future revision of this standard.
-
-   For array libraries which concern themselves with multi-device support, including CPU and GPU,
-   they are free to expose a library-specific device object (e.g., for creating an
-   array on a particular device). While a library-specific device object can be used as input to
-   ``to_device``, beware that this will mean non-portability as code will be specific to
-   that library.
+   The current API standard does **not** include a universal ``Device`` object
+   recognized by all compliant libraries. Accordingly, the standard does not
+   provide a means of instantiating a ``Device`` object to point to a specific
+   physical or logical device.
+
+   The choice to not include a standardized ``Device`` object may be revisited
+   in a future revision of this standard.
+
+   For array libraries which concern themselves with multi-device support,
+   including CPU and GPU, they are free to expose a library-specific device
+   object (e.g., for creating an array on a particular device). While a
+   library-specific device object can be used as input to ``to_device``, beware
+   that this will mean non-portability as code will be specific to that
+   library.
 
 Semantics
 ---------
diff --git a/src/_array_api_conf.py b/src/_array_api_conf.py
@@ -72,6 +72,10 @@
     ("py:class", ".*PyCapsule"),
     ("py:class", ".*finfo_object"),
     ("py:class", ".*iinfo_object"),
+    ("py:class", ".*Info"),
+    ("py:class", ".*Capabilities"),
+    ("py:class", ".*DefaultDataTypes"),
+    ("py:class", ".*DataTypes"),
 ]
 # In array_object.py we have to use aliased names for some types because they
 # would otherwise refer back to method objects of array
diff --git a/src/array_api_stubs/_draft/__init__.py b/src/array_api_stubs/_draft/__init__.py
@@ -16,6 +16,7 @@
 from .utility_functions import *
 from . import linalg
 from . import fft
+from . import info
 
 
 __array_api_version__: str = "YYYY.MM"
diff --git a/src/array_api_stubs/_draft/_types.py b/src/array_api_stubs/_draft/_types.py
@@ -25,6 +25,10 @@
     "finfo_object",
     "iinfo_object",
     "Enum",
+    "DefaultDataTypes",
+    "DataTypes",
+    "Capabilities",
+    "Info",
 ]
 
 from dataclasses import dataclass
@@ -35,6 +39,7 @@
     Optional,
     Sequence,
     Tuple,
+    TypedDict,
     TypeVar,
     Union,
     Protocol,
@@ -83,3 +88,57 @@ def __getitem__(self, key: int, /) -> Union[_T_co, NestedSequence[_T_co]]:
 
     def __len__(self, /) -> int:
         ...
+
+
+class Info(Protocol):
+    """Namespace returned by `__array_namespace_info__`."""
+
+    def capabilities(self) -> Capabilities:
+        ...
+
+    def default_device(self) -> device:
+        ...
+
+    def default_dtypes(self, *, device: Optional[device]) -> DefaultDataTypes:
+        ...
+
+    def devices(self) -> List[device]:
+        ...
+
+    def dtypes(
+        self, *, device: Optional[device], kind: Optional[Union[str, Tuple[str, ...]]]
+    ) -> DataTypes:
+        ...
+
+
+DefaultDataTypes = TypedDict(
+    "DefaultDataTypes",
+    {
+        "real floating": dtype,
+        "complex floating": dtype,
+        "integral": dtype,
+        "indexing": dtype,
+    },
+)
+DataTypes = TypedDict(
+    "DataTypes",
+    {
+        "bool": dtype,
+        "float32": dtype,
+        "float64": dtype,
+        "complex64": dtype,
+        "complex128": dtype,
+        "int8": dtype,
+        "int16": dtype,
+        "int32": dtype,
+        "int64": dtype,
+        "uint8": dtype,
+        "uint16": dtype,
+        "uint32": dtype,
+        "uint64": dtype,
+    },
+    total=False,
+)
+Capabilities = TypedDict(
+    "Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool}
+)
diff --git a/src/array_api_stubs/_draft/info.py b/src/array_api_stubs/_draft/info.py
@@ -0,0 +1,170 @@
+__all__ = [
+    "__array_namespace_info__",
+    "capabilities",
+    "default_device",
+    "default_dtypes",
+    "devices",
+    "dtypes",
+]
+
+from ._types import (
+    Optional,
+    Union,
+    Tuple,
+    List,
+    device,
+    dtype,
+    DefaultDataTypes,
+    DataTypes,
+    Capabilities,
+    Info,
+)
+
+
+def __array_namespace_info__() -> Info:
+    """
+    Returns a namespace with Array API namespace inspection utilities.
+
+    Returns
+    -------
+    out: Info
+        An object containing Array API namespace inspection utilities.
+
+    Notes
+    -----
+
+    The returned object may be either a namespace or a class, so long as an Array API user can access inspection utilities as follows:
+
+    ::
+
+      info = xp.__array_namespace_info__()
+      info.capabilities()
+      info.devices()
+      info.dtypes()
+      info.default_dtypes()
+      # ...
+    """
+
+
+def capabilities() -> Capabilities:
+    """
+    Returns a dictionary of array library capabilities.
+
+    The dictionary must contain the following keys:
+
+    -   `"boolean indexing"`: boolean indicating whether an array library supports boolean indexing. If a conforming implementation fully supports boolean indexing in compliance with this specification (see :ref:`indexing`), the corresponding dictionary value must be ``True``; otherwise, the value must be ``False``.
+    -   `"data-dependent shapes"`: boolean indicating whether an array library supports data-dependent output shapes. If a conforming implementation fully supports all APIs included in this specification (excluding boolean indexing) which have data-dependent output shapes, as explicitly demarcated throughout the specification, the corresponding dictionary value must be ``True``; otherwise, the value must be ``False``.
+
+    Returns
+    -------
+    out: Capabilities
+        a dictionary of array library capabilities.
+    """
+
+
+def default_device() -> device:
+    """
+    Returns the default device.
+
+    Returns
+    -------
+    out: device
+        an object corresponding to the default device.
+    """
+
+
+def default_dtypes(
+    *,
+    device: Optional[device] = None,
+) -> DefaultDataTypes:
+    """
+    Returns a dictionary containing default data types.
+
+    The dictionary must have the following keys:
+
+    -   `"real floating"`: default real floating-point data type.
+    -   `"complex floating"`: default complex floating-point data type.
+    -   `"integral"`: default integral data type.
+    -   `"indexing"`: default array index data type.
+
+    Dictionary values must be the corresponding data type object.
+
+    Parameters
+    ----------
+    device: Optional[device]
+        device for which to return default data types. If ``device`` is ``None``, the returned data types must be the default data types for the current device; otherwise, the returned data types must be default data types specific to the specified device. Default: ``None``.
+
+        .. note::
+           Some array libraries have the concept of a device context manager, allowing library consumers to manage the current device context. When ``device`` is ``None``, libraries supporting a device context should return the default data types for the current device. For libraries without a context manager or supporting only a single device, those libraries should return the default data types for the default device.
+
+    Returns
+    -------
+    out: DefaultDataTypes
+        a dictionary containing the default data type for respective data type kinds.
+    """
+
+
+def dtypes(
+    *,
+    device: Optional[device] = None,
+    kind: Optional[Union[str, Tuple[str, ...]]] = None,
+) -> DataTypes:
+    """
+    Returns a dictionary of supported *Array API* data types.
+
+    .. note::
+       While specification-conforming array libraries may support additional data types which are not present in this specification, data types which are not present in this specification should not be included in the returned dictionary.
+
+    .. note::
+       Specification-conforming array libraries must only return supported data types having expected properties as described in :ref:`data-types`. For example, if a library decides to alias ``float32`` as ``float64``, that library must not include ``float64`` in the dictionary of supported data types.
+
+    Parameters
+    ----------
+    kind: Optional[Union[str, Tuple[str, ...]]]
+        data type kind.
+
+        -   If ``kind`` is ``None``, the function must return a dictionary containing all supported Array API data types.
+
+        -   If ``kind`` is a string, the function must return a dictionary containing the data types belonging to the specified data type kind. The following data type kinds must be supported:
+
+            -   ``'bool'``: boolean data types (e.g., ``bool``).
+            -   ``'signed integer'``: signed integer data types (e.g., ``int8``, ``int16``, ``int32``, ``int64``).
+            -   ``'unsigned integer'``: unsigned integer data types (e.g., ``uint8``, ``uint16``, ``uint32``, ``uint64``).
+            -   ``'integral'``: integer data types. Shorthand for ``('signed integer', 'unsigned integer')``.
+            -   ``'real floating'``: real-valued floating-point data types (e.g., ``float32``, ``float64``).
+            -   ``'complex floating'``: complex floating-point data types (e.g., ``complex64``, ``complex128``).
+            -   ``'numeric'``: numeric data types. Shorthand for ``('integral', 'real floating', 'complex floating')``.
+
+        -   If ``kind`` is a tuple, the tuple specifies a union of data type kinds, and the function must return a dictionary containing the data types belonging to at least one of the specified data type kinds.
+
+        Default: ``None``.
+    device: Optional[device]
+        device for which to return supported data types. If ``device`` is ``None``, the returned data types must be the supported data types for the current device; otherwise, the returned data types must be supported data types specific to the specified device. Default: ``None``.
+
+        .. note::
+           Some array libraries have the concept of a device context manager, allowing library consumers to manage the current device context. When ``device`` is ``None``, libraries supporting a device context should return the supported data types for the current device. For libraries without a context manager or supporting only a single device, those libraries should return the supported data types for the default device.
+
+    Returns
+    -------
+    out: DataTypes
+        a dictionary containing supported data types.
+
+        .. note::
+           Dictionary keys must only consist of canonical names as defined in :ref:`data-types`.
+    """
+
+
+def devices() -> List[device]:
+    """
+    Returns a list of supported devices which are available at runtime.
+
+    Returns
+    -------
+    out: List[device]
+        a list of supported devices.
+
+    Notes
+    -----
+
+    Each device object (see :ref:`device-support`) in the list of returned devices must be an object which can be provided as a valid keyword-argument to array creation functions.
+    """
