add DeviceArray type alias for typing

Author: alexfikl

arraycontext/__init__.py

@@ -29,7 +29,7 @@
 """
 
 import sys
-from .context import ArrayContext
+from .context import ArrayContext, DeviceArray, DeviceScalar
 
 from .transform_metadata import (CommonSubexpressionTag,
         ElementwiseMapKernelTag)
@@ -74,7 +74,7 @@
 
 
 __all__ = (
-        "ArrayContext",
+        "ArrayContext", "DeviceScalar", "DeviceArray",
 
         "CommonSubexpressionTag",
         "ElementwiseMapKernelTag",

arraycontext/container/traversal.py

@@ -67,7 +67,7 @@
 
 import numpy as np
 
-from arraycontext.context import ArrayContext
+from arraycontext.context import ArrayContext, DeviceArray
 from arraycontext.container import (
         ContainerT, ArrayOrContainerT, NotAnArrayContainerError,
         serialize_container, deserialize_container)
@@ -355,7 +355,7 @@ def rec(keys: Tuple[Union[str, int], ...],
 def map_reduce_array_container(
         reduce_func: Callable[[Iterable[Any]], Any],
         map_func: Callable[[Any], Any],
-        ary: ArrayOrContainerT) -> Any:
+        ary: ArrayOrContainerT) -> "DeviceArray":
     """Perform a map-reduce over array containers.
 
     :param reduce_func: callable used to reduce over the components of *ary*
@@ -378,7 +378,7 @@ def map_reduce_array_container(
 def multimap_reduce_array_container(
         reduce_func: Callable[[Iterable[Any]], Any],
         map_func: Callable[..., Any],
-        *args: Any) -> Any:
+        *args: Any) -> "DeviceArray":
     r"""Perform a map-reduce over multiple array containers.
 
     :param reduce_func: callable used to reduce over the components of any
@@ -401,7 +401,7 @@ def _reduce_wrapper(ary: ContainerT, iterable: Iterable[Tuple[Any, Any]]) -> Any
 def rec_map_reduce_array_container(
         reduce_func: Callable[[Iterable[Any]], Any],
         map_func: Callable[[Any], Any],
-        ary: ArrayOrContainerT) -> Any:
+        ary: ArrayOrContainerT) -> "DeviceArray":
     """Perform a map-reduce over array containers recursively.
 
     :param reduce_func: callable used to reduce over the components of *ary*
@@ -455,7 +455,7 @@ def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
 def rec_multimap_reduce_array_container(
         reduce_func: Callable[[Iterable[Any]], Any],
         map_func: Callable[..., Any],
-        *args: Any) -> Any:
+        *args: Any) -> "DeviceArray":
     r"""Perform a map-reduce over multiple array containers recursively.
 
     :param reduce_func: callable used to reduce over the components of any

arraycontext/context.py

@@ -74,6 +74,19 @@
 ---------------------------------
 
 .. currentmodule:: arraycontext
+
+.. class:: DeviceArray
+
+    A (type alias for an) array type supported by the :class:`ArrayContext`
+    meant to aid in typing annotations. For a explicit list of supported types
+    see :attr:`ArrayContext.array_types`.
+
+.. class:: DeviceScalar
+
+    A (type alias for a) scalar type supported by the :class:`ArrayContext`
+    meant to aid in typing annotations, e.g. for reductions. In :mod:`numpy`
+    terminology, this is just an array with a shape of ``()``.
+
 .. autoclass:: ArrayContext
 """
 
@@ -110,6 +123,10 @@
 from pytools.tag import Tag
 
 
+DeviceArray = Any
+DeviceScalar = Any
+
+
 # {{{ ArrayContext
 
 class ArrayContext(ABC):

doc/conf.py

@@ -14,6 +14,11 @@
 version = ".".join(str(x) for x in ver_dic["VERSION"])
 release = ver_dic["VERSION_TEXT"]
 
+autodoc_type_aliases = {
+        "DeviceScalar": "arraycontext.DeviceScalar",
+        "DeviceArray": "arraycontext.DeviceArray",
+        }
+
 intersphinx_mapping = {
     "https://docs.python.org/3/": None,
     "https://numpy.org/doc/stable/": None,