Merge branch 'main' into cupyactx

Author: matthiasdiener

.basedpyright/baseline.json

@@ -31,25 +31,6 @@ jobs:
                 pip install ruff
                 ruff check
 
-    mypy:
-        name: Mypy
-        runs-on: ubuntu-latest
-        steps:
-        -   uses: actions/checkout@v4
-        -   name: "Main Script"
-            run: |
-                set -x
-                USE_CONDA_BUILD=1
-                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
-                echo "- cupy" >> "$CONDA_ENVIRONMENT"
-
-                curl -L -O https://tiker.net/ci-support-v0
-                . ./ci-support-v0
-
-                build_py_project_in_conda_env
-                python -m pip install mypy pytest
-                ./run-mypy.sh
-
     basedpyright:
         runs-on: ubuntu-latest
         steps:

.github/workflows/ci.yml

@@ -115,20 +115,6 @@ Ruff:
   except:
   - tags
 
-Mypy:
-  script: |
-    EXTRA_INSTALL="mypy pytest"
-
-    curl -L -O https://tiker.net/ci-support-v0
-    . ./ci-support-v0
-
-    build_py_project_in_venv
-    ./run-mypy.sh
-  tags:
-  - python3
-  except:
-  - tags
-
 Downstream:
   parallel:
     matrix:

.gitlab-ci.yml

@@ -77,6 +77,8 @@
 
 import numpy as np
 
+from pymbolic.typing import Integer
+
 from arraycontext.container import (
     ArrayContainer,
     NotAnArrayContainerError,
@@ -91,7 +93,6 @@
     ArrayOrContainer,
     ArrayOrContainerOrScalar,
     ArrayOrContainerT,
-    ArrayT,
     ScalarLike,
 )
 
@@ -400,21 +401,20 @@ def keyed_map_array_container(
 
 
 def rec_keyed_map_array_container(
-        f: Callable[[tuple[SerializationKey, ...], ArrayT], ArrayT],
+        f: Callable[[tuple[SerializationKey, ...], Array], Array],
         ary: ArrayOrContainer) -> ArrayOrContainer:
     """
     Works similarly to :func:`rec_map_array_container`, except that *f* also
     takes in a traversal path to the leaf array. The traversal path argument is
     passed in as a tuple of identifiers of the arrays traversed before reaching
     the current array.
     """
-
     def rec(keys: tuple[SerializationKey, ...],
-            ary_: ArrayOrContainerT) -> ArrayOrContainerT:
+            ary_: ArrayOrContainer) -> ArrayOrContainer:
         try:
             iterable = serialize_container(ary_)
         except NotAnArrayContainerError:
-            return cast(ArrayOrContainerT, f(keys, cast(ArrayT, ary_)))
+            return cast(ArrayOrContainer, f(keys, cast(Array, ary_)))
         else:
             return deserialize_container(ary_, [
                 (key, rec((*keys, key), subary)) for key, subary in iterable
@@ -777,7 +777,7 @@ def unflatten(
         checks are skipped.
     """
     # NOTE: https://github.com/python/mypy/issues/7057
-    offset = 0
+    offset: int = 0
     common_dtype = None
 
     def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
@@ -790,7 +790,11 @@ def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
 
             # {{{ validate subary
 
-            if (offset + template_subary_c.size) > ary.size:
+            if (
+                    isinstance(offset, Integer)
+                    and isinstance(template_subary_c.size, Integer)
+                    and isinstance(ary.size, Integer)
+                    and (offset + template_subary_c.size) > ary.size):
                 raise ValueError("'template' and 'ary' sizes do not match: "
                     "'template' is too large") from None
 
@@ -813,6 +817,12 @@ def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
 
             # {{{ reshape
 
+            if not isinstance(template_subary_c.size, Integer):
+                raise NotImplementedError(
+                    "unflatten is not implemented for arrays with array-valued "
+                    "size.") from None
+
+            # FIXME: Not sure how to make the slicing part work for Array-valued sizes
             flat_subary = ary[offset:offset + template_subary_c.size]
             try:
                 subary = actx.np.reshape(flat_subary,
@@ -871,15 +881,15 @@ def _unflatten(template_subary: ArrayOrContainer) -> ArrayOrContainer:
 
 
 def flat_size_and_dtype(
-        ary: ArrayOrContainer) -> tuple[int, np.dtype[Any] | None]:
+        ary: ArrayOrContainer) -> tuple[Array | Integer, np.dtype[Any] | None]:
     """
     :returns: a tuple ``(size, dtype)`` that would be the length and
         :class:`numpy.dtype` of the one-dimensional array returned by
         :func:`flatten`.
     """
     common_dtype = None
 
-    def _flat_size(subary: ArrayOrContainer) -> int:
+    def _flat_size(subary: ArrayOrContainer) -> Array | Integer:
         nonlocal common_dtype
 
         try:

arraycontext/container/traversal.py

@@ -174,6 +174,7 @@
 import numpy as np
 from typing_extensions import Self
 
+from pymbolic.typing import Integer, Scalar as _Scalar
 from pytools import memoize_method
 from pytools.tag import ToTagSetConvertible
 
@@ -202,11 +203,11 @@ class Array(Protocol):
     """
 
     @property
-    def shape(self) -> tuple[int, ...]:
+    def shape(self) -> tuple[Array | Integer, ...]:
         ...
 
     @property
-    def size(self) -> int:
+    def size(self) -> Array | Integer:
         ...
 
     @property
@@ -220,22 +221,27 @@ def dtype(self) -> np.dtype[Any]:
     def __getitem__(self, index: Any) -> Array:
         ...
 
-    # some basic arithmetic that's supposed to work
-    def __neg__(self) -> Self: ...
-    def __abs__(self) -> Self: ...
-    def __add__(self, other: Self | ScalarLike) -> Self: ...
-    def __radd__(self, other: Self | ScalarLike) -> Self: ...
-    def __sub__(self, other: Self | ScalarLike) -> Self: ...
-    def __rsub__(self, other: Self | ScalarLike) -> Self: ...
-    def __mul__(self, other: Self | ScalarLike) -> Self: ...
-    def __rmul__(self, other: Self | ScalarLike) -> Self: ...
-    def __truediv__(self, other: Self | ScalarLike) -> Self: ...
-    def __rtruediv__(self, other: Self | ScalarLike) -> Self: ...
+    # Some basic arithmetic that's supposed to work
+    # Need to return Array instead of Self because for some array types, arithmetic
+    # operations on one subtype may result in a different subtype.
+    # For example, pytato arrays: <Placeholder> + 1 -> <IndexLambda>
+    def __neg__(self) -> Array: ...
+    def __abs__(self) -> Array: ...
+    def __add__(self, other: Self | ScalarLike) -> Array: ...
+    def __radd__(self, other: Self | ScalarLike) -> Array: ...
+    def __sub__(self, other: Self | ScalarLike) -> Array: ...
+    def __rsub__(self, other: Self | ScalarLike) -> Array: ...
+    def __mul__(self, other: Self | ScalarLike) -> Array: ...
+    def __rmul__(self, other: Self | ScalarLike) -> Array: ...
+    def __pow__(self, other: Self | ScalarLike) -> Array: ...
+    def __rpow__(self, other: Self | ScalarLike) -> Array: ...
+    def __truediv__(self, other: Self | ScalarLike) -> Array: ...
+    def __rtruediv__(self, other: Self | ScalarLike) -> Array: ...
 
 
 # deprecated, use ScalarLike instead
-ScalarLike: TypeAlias = int | float | complex | np.generic
-Scalar = ScalarLike
+Scalar = _Scalar
+ScalarLike = Scalar
 ScalarLikeT = TypeVar("ScalarLikeT", bound=ScalarLike)
 
 # NOTE: I'm kind of not sure about the *Tc versions of these type variables.

arraycontext/context.py

@@ -32,7 +32,7 @@
 """
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 from warnings import warn
 
 import numpy as np
@@ -51,7 +51,8 @@
 
 if TYPE_CHECKING:
     import loopy as lp
-    import pyopencl
+    import pyopencl as cl
+    import pyopencl.array as cl_array
 
 
 # {{{ PyOpenCLArrayContext
@@ -81,9 +82,17 @@ class PyOpenCLArrayContext(ArrayContext):
     .. automethod:: transform_loopy_program
     """
 
+    context: cl.Context
+    queue: cl.CommandQueue
+    allocator: cl_array.Allocator | None
+
+    _force_device_scalars: Literal[True]
+    _passed_force_device_scalars: bool
+    _wait_event_queue_length: int
+
     def __init__(self,
-            queue: pyopencl.CommandQueue,
-            allocator: pyopencl.tools.AllocatorBase | None = None,
+            queue: cl.CommandQueue,
+            allocator: cl_array.Allocator | None = None,
             wait_event_queue_length: int | None = None,
             force_device_scalars: bool | None = None) -> None:
         r"""

arraycontext/impl/pyopencl/__init__.py

@@ -41,7 +41,7 @@
     rec_multimap_array_container,
     rec_multimap_reduce_array_container,
 )
-from arraycontext.context import Array, ArrayOrContainer
+from arraycontext.context import Array as actx_Array, ArrayOrContainer
 from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace
 from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
 from arraycontext.loopy import LoopyBasedFakeNumpyNamespace
@@ -206,7 +206,7 @@ def _any(ary):
                 _any,
                 a)
 
-    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> actx_Array:
         actx = self._array_context
         queue = actx.queue
 

arraycontext/impl/pyopencl/fake_numpy.py

@@ -7,10 +7,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
+from numpy.typing import DTypeLike
 
+import pyopencl as cl
 import pyopencl.array as cla
 from pytools import memoize
 from pytools.tag import Tag, Taggable, ToTagSetConvertible
@@ -74,6 +76,9 @@ class TaggableCLArray(cla.Array, Taggable):
         record application-specific metadata to drive the optimizations in
         :meth:`arraycontext.PyOpenCLArrayContext.transform_loopy_program`.
     """
+    tags: frozenset[Tag]
+    axes: tuple[Axis, ...]
+
     def __init__(self, cq, shape, dtype, order="C", allocator=None,
                  data=None, offset=0, strides=None, events=None, _flags=None,
                  _fast=False, _size=None, _context=None, _queue=None,
@@ -165,13 +170,20 @@ def to_tagged_cl_array(ary: cla.Array,
 # }}}
 
 
+_EMPTY_TAG_SET: frozenset[Tag] = frozenset()
+
+
 # {{{ creation
 
-def empty(queue, shape, dtype=float, *,
-        axes: tuple[Axis, ...] | None = None,
-        tags: frozenset[Tag] = frozenset(),
-        order: str = "C",
-        allocator=None) -> TaggableCLArray:
+def empty(
+            queue: cl.CommandQueue,
+            shape: tuple[int, ...] | int,
+            dtype: DTypeLike = float,
+            *, axes: tuple[Axis, ...] | None = None,
+            tags: frozenset[Tag] = _EMPTY_TAG_SET,
+            order: Literal["C"] | Literal["F"] = "C",
+            allocator: cla.Allocator | None = None,
+        ) -> TaggableCLArray:
     if dtype is not None:
         dtype = np.dtype(dtype)
 
@@ -181,11 +193,15 @@ def empty(queue, shape, dtype=float, *,
         order=order, allocator=allocator)
 
 
-def zeros(queue, shape, dtype=float, *,
-        axes: tuple[Axis, ...] | None = None,
-        tags: frozenset[Tag] = frozenset(),
-        order: str = "C",
-        allocator=None) -> TaggableCLArray:
+def zeros(
+            queue: cl.CommandQueue,
+            shape: tuple[int, ...] | int,
+            dtype: DTypeLike = float,
+            *, axes: tuple[Axis, ...] | None = None,
+            tags: frozenset[Tag] = _EMPTY_TAG_SET,
+            order: Literal["C"] | Literal["F"] = "C",
+            allocator: cla.Allocator | None = None,
+        ) -> TaggableCLArray:
     result = empty(
         queue, shape, dtype=dtype, axes=axes, tags=tags,
         order=order, allocator=allocator)
@@ -194,10 +210,13 @@ def zeros(queue, shape, dtype=float, *,
     return result
 
 
-def to_device(queue, ary, *,
-        axes: tuple[Axis, ...] | None = None,
-        tags: frozenset[Tag] = frozenset(),
-        allocator=None):
+def to_device(
+            queue: cl.CommandQueue,
+            ary: np.ndarray[Any],
+            *, axes: tuple[Axis, ...] | None = None,
+            tags: frozenset[Tag] = _EMPTY_TAG_SET,
+            allocator: cla.Allocator | None = None,
+        ) -> TaggableCLArray:
     return to_tagged_cl_array(
         cla.to_device(queue, ary, allocator=allocator),
         axes=axes, tags=tags)