From 7ab5211bd96c0121c87854e8f8365d282cf85633 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 16 Feb 2024 16:40:22 -0600
Subject: [PATCH 01/34] Implement CupyArrayContext

Co-authored-by: Kaushik Kulkarni <kaushikcfd@gmail.com>
---
 .github/workflows/ci.yml             |   1 +
 arraycontext/__init__.py             |   5 +-
 arraycontext/impl/cupy/__init__.py   | 136 +++++++++++++++++++++++
 arraycontext/impl/cupy/fake_numpy.py | 156 +++++++++++++++++++++++++++
 arraycontext/pytest.py               |  18 ++++
 test/test_arraycontext.py            |  49 +++++++--
 6 files changed, 354 insertions(+), 11 deletions(-)
 create mode 100644 arraycontext/impl/cupy/__init__.py
 create mode 100644 arraycontext/impl/cupy/fake_numpy.py

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 29952752..4b191f96 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -49,6 +49,7 @@ jobs:
                 . ./ci-support-v0
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
+                conda install cupy
                 ./run-mypy.sh
 
     pytest3_pocl:
diff --git a/arraycontext/__init__.py b/arraycontext/__init__.py
index b01b9917..a0a9deb6 100644
--- a/arraycontext/__init__.py
+++ b/arraycontext/__init__.py
@@ -50,6 +50,7 @@
     Array, ArrayContext, ArrayOrContainer, ArrayOrContainerOrScalar,
     ArrayOrContainerOrScalarT, ArrayOrContainerT, ArrayT, Scalar, ScalarLike,
     tag_axes)
+from .impl.cupy import CupyArrayContext
 from .impl.jax import EagerJAXArrayContext
 from .impl.pyopencl import PyOpenCLArrayContext
 from .impl.pytato import PytatoJAXArrayContext, PytatoPyOpenCLArrayContext
@@ -106,7 +107,9 @@
         "PytestArrayContextFactory",
         "PytestPyOpenCLArrayContextFactory",
         "pytest_generate_tests_for_array_contexts",
-        "pytest_generate_tests_for_pyopencl_array_context"
+        "pytest_generate_tests_for_pyopencl_array_context",
+
+        "CupyArrayContext",
         )
 
 
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
new file mode 100644
index 00000000..f58d16d6
--- /dev/null
+++ b/arraycontext/impl/cupy/__init__.py
@@ -0,0 +1,136 @@
+"""
+.. currentmodule:: arraycontext
+
+
+A mod :`cupy`-based array context.
+
+.. autoclass:: CupyArrayContext
+"""
+__copyright__ = """
+Copyright (C) 2024 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from collections.abc import Mapping
+
+
+try:
+    import cupy as cp  # type: ignore[import-untyped]
+except ModuleNotFoundError:
+    pass
+
+import loopy as lp
+
+from arraycontext.container.traversal import (
+    rec_map_array_container, with_array_context)
+from arraycontext.context import ArrayContext
+
+
+class CupyArrayContext(ArrayContext):
+    """
+    A :class:`ArrayContext` that uses :mod:`cupy.ndarray` to represent arrays
+
+
+    .. automethod:: __init__
+    """
+    def __init__(self):
+        super().__init__()
+        self._loopy_transform_cache: \
+                Mapping["lp.TranslationUnit", "lp.TranslationUnit"] = {}
+
+        self.array_types = (cp.ndarray,)
+
+    def _get_fake_numpy_namespace(self):
+        from .fake_numpy import CupyFakeNumpyNamespace
+        return CupyFakeNumpyNamespace(self)
+
+    # {{{ ArrayContext interface
+
+    def clone(self):
+        return type(self)()
+
+    def empty(self, shape, dtype):
+        return cp.empty(shape, dtype=dtype)
+
+    def zeros(self, shape, dtype):
+        return cp.zeros(shape, dtype)
+
+    def from_numpy(self, np_array):
+        return cp.array(np_array)
+
+    def to_numpy(self, array):
+        return cp.asnumpy(array)
+
+    def call_loopy(self, t_unit, **kwargs):
+        t_unit = t_unit.copy(target=lp.ExecutableCTarget())
+        try:
+            t_unit = self._loopy_transform_cache[t_unit]
+        except KeyError:
+            orig_t_unit = t_unit
+            t_unit = self.transform_loopy_program(t_unit)
+            self._loopy_transform_cache[orig_t_unit] = t_unit
+            del orig_t_unit
+
+        _, result = t_unit(**kwargs)
+
+        return result
+
+    def freeze(self, array):
+        def _freeze(ary):
+            return cp.asnumpy(ary)
+
+        return with_array_context(rec_map_array_container(_freeze, array), actx=None)
+
+    def thaw(self, array):
+        def _thaw(ary):
+            return cp.array(ary)
+
+        return with_array_context(rec_map_array_container(_thaw, array), actx=self)
+
+    # }}}
+
+    def transform_loopy_program(self, t_unit):
+        raise ValueError("CupyArrayContext does not implement "
+                         "transform_loopy_program. Sub-classes are supposed "
+                         "to implement it.")
+
+    def tag(self, tags, array):
+        # No tagging support in CupyArrayContext
+        return array
+
+    def tag_axis(self, iaxis, tags, array):
+        return array
+
+    def einsum(self, spec, *args, arg_names=None, tagged=()):
+        return cp.einsum(spec, *args)
+
+    @property
+    def permits_inplace_modification(self):
+        return True
+
+    @property
+    def supports_nonscalar_broadcasting(self):
+        return True
+
+    @property
+    def permits_advanced_indexing(self):
+        return True
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
new file mode 100644
index 00000000..3e6fa805
--- /dev/null
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -0,0 +1,156 @@
+__copyright__ = """
+Copyright (C) 2024 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+from functools import partial, reduce
+
+import cupy as cp  # type: ignore[import-untyped]  # pylint: disable=import-error
+
+from arraycontext.container import is_array_container
+from arraycontext.container.traversal import (
+    multimap_reduce_array_container, rec_map_array_container,
+    rec_map_reduce_array_container, rec_multimap_array_container,
+    rec_multimap_reduce_array_container)
+from arraycontext.fake_numpy import (
+    BaseFakeNumpyLinalgNamespace, BaseFakeNumpyNamespace)
+
+
+class CupyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
+    # Everything is implemented in the base class for now.
+    pass
+
+
+_NUMPY_UFUNCS = {"abs", "sin", "cos", "tan", "arcsin", "arccos", "arctan",
+                 "sinh", "cosh", "tanh", "exp", "log", "log10", "isnan",
+                 "sqrt", "concatenate", "transpose",
+                 "ones_like", "maximum", "minimum", "where", "conj", "arctan2",
+                 }
+
+
+class CupyFakeNumpyNamespace(BaseFakeNumpyNamespace):
+    """
+    A :mod:`numpy` mimic for :class:`CupyArrayContext`.
+    """
+    def _get_fake_numpy_linalg_namespace(self):
+        return CupyFakeNumpyLinalgNamespace(self._array_context)
+
+    def __getattr__(self, name):
+
+        if name in _NUMPY_UFUNCS:
+            from functools import partial
+            return partial(rec_multimap_array_container,
+                           getattr(cp, name))
+
+        raise NotImplementedError
+
+    def sum(self, a, axis=None, dtype=None):
+        return rec_map_reduce_array_container(sum, partial(cp.sum,
+                                                           axis=axis,
+                                                           dtype=dtype),
+                                              a)
+
+    def min(self, a, axis=None):
+        return rec_map_reduce_array_container(
+                partial(reduce, cp.minimum), partial(cp.amin, axis=axis), a)
+
+    def max(self, a, axis=None):
+        return rec_map_reduce_array_container(
+                partial(reduce, cp.maximum), partial(cp.amax, axis=axis), a)
+
+    def stack(self, arrays, axis=0):
+        return rec_multimap_array_container(
+                lambda *args: cp.stack(args, axis=axis),
+                *arrays)
+
+    def broadcast_to(self, array, shape):
+        return rec_map_array_container(partial(cp.broadcast_to, shape=shape), array)
+
+    # {{{ relational operators
+
+    def equal(self, x, y):
+        return rec_multimap_array_container(cp.equal, x, y)
+
+    def not_equal(self, x, y):
+        return rec_multimap_array_container(cp.not_equal, x, y)
+
+    def greater(self, x, y):
+        return rec_multimap_array_container(cp.greater, x, y)
+
+    def greater_equal(self, x, y):
+        return rec_multimap_array_container(cp.greater_equal, x, y)
+
+    def less(self, x, y):
+        return rec_multimap_array_container(cp.less, x, y)
+
+    def less_equal(self, x, y):
+        return rec_multimap_array_container(cp.less_equal, x, y)
+
+    # }}}
+
+    def ravel(self, a, order="C"):
+        return rec_map_array_container(partial(cp.ravel, order=order), a)
+
+    def vdot(self, x, y, dtype=None):
+        if dtype is not None:
+            raise NotImplementedError("only 'dtype=None' supported.")
+
+        return rec_multimap_reduce_array_container(sum, cp.vdot, x, y)
+
+    def any(self, a):
+        return rec_map_reduce_array_container(partial(reduce, cp.logical_or),
+                                              lambda subary: cp.any(subary), a)
+
+    def all(self, a):
+        return rec_map_reduce_array_container(partial(reduce, cp.logical_and),
+                                              lambda subary: cp.all(subary), a)
+
+    def array_equal(self, a, b):
+        if type(a) is not type(b):
+            return False
+        elif not is_array_container(a):
+            if a.shape != b.shape:
+                return False
+            else:
+                return cp.all(cp.equal(a, b))
+        else:
+            try:
+                return multimap_reduce_array_container(partial(reduce,
+                                                           cp.logical_and),
+                                                   self.array_equal, a, b)
+            except TypeError:
+                return True
+
+    def zeros_like(self, ary):
+        return rec_multimap_array_container(cp.zeros_like, ary)
+
+    def reshape(self, a, newshape, order="C"):
+        return rec_map_array_container(
+                lambda ary: ary.reshape(newshape, order=order),
+                a)
+
+    def arange(self, *args, **kwargs):
+        return cp.arange(*args, **kwargs)
+
+    def linspace(self, *args, **kwargs):
+        return cp.linspace(*args, **kwargs)
+
+# vim: fdm=marker
diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py
index 4fce5885..f8de488a 100644
--- a/arraycontext/pytest.py
+++ b/arraycontext/pytest.py
@@ -224,6 +224,23 @@ def __str__(self):
         return "<PytatoJAXArrayContext>"
 
 
+class _PytestCupyArrayContextFactory(PytestArrayContextFactory):
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            import cupy  # type: ignore[import-untyped]  # noqa: F401
+            return True
+        except ImportError:
+            return False
+
+    def __call__(self):
+        from arraycontext import CupyArrayContext
+        return CupyArrayContext()
+
+    def __str__(self):
+        return "<CupyArrayContext>"
+
+
 _ARRAY_CONTEXT_FACTORY_REGISTRY: \
         Dict[str, Type[PytestArrayContextFactory]] = {
                 "pyopencl": _PytestPyOpenCLArrayContextFactoryWithClass,
@@ -232,6 +249,7 @@ def __str__(self):
                 "pytato:pyopencl": _PytestPytatoPyOpenCLArrayContextFactory,
                 "pytato:jax": _PytestPytatoJaxArrayContextFactory,
                 "eagerjax": _PytestEagerJaxArrayContextFactory,
+                "cupy": _PytestCupyArrayContextFactory,
                 }
 
 
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index e53f4295..712d5f24 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -30,13 +30,15 @@
 from pytools.obj_array import make_obj_array
 
 from arraycontext import (  # noqa: F401
-    ArrayContainer, ArrayContext, EagerJAXArrayContext, FirstAxisIsElementsTag,
-    PyOpenCLArrayContext, PytatoPyOpenCLArrayContext, dataclass_array_container,
-    deserialize_container, pytest_generate_tests_for_array_contexts,
-    serialize_container, tag_axes, with_array_context, with_container_arithmetic)
+    ArrayContainer, ArrayContext, CupyArrayContext, EagerJAXArrayContext,
+    FirstAxisIsElementsTag, PyOpenCLArrayContext, PytatoPyOpenCLArrayContext,
+    dataclass_array_container, deserialize_container,
+    pytest_generate_tests_for_array_contexts, serialize_container, tag_axes,
+    with_array_context, with_container_arithmetic)
 from arraycontext.pytest import (
-    _PytestEagerJaxArrayContextFactory, _PytestPyOpenCLArrayContextFactoryWithClass,
-    _PytestPytatoJaxArrayContextFactory, _PytestPytatoPyOpenCLArrayContextFactory)
+    _PytestCupyArrayContextFactory, _PytestEagerJaxArrayContextFactory,
+    _PytestPyOpenCLArrayContextFactoryWithClass, _PytestPytatoJaxArrayContextFactory,
+    _PytestPytatoPyOpenCLArrayContextFactory)
 
 
 logger = logging.getLogger(__name__)
@@ -84,6 +86,7 @@ class _PytatoPyOpenCLArrayContextForTestsFactory(
     _PytatoPyOpenCLArrayContextForTestsFactory,
     _PytestEagerJaxArrayContextFactory,
     _PytestPytatoJaxArrayContextFactory,
+    _PytestCupyArrayContextFactory,
     ])
 
 
@@ -412,6 +415,11 @@ def test_array_context_np_like(actx_factory, sym_name, n_args, dtype):
             actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
 
     for c in (42.0,) + _get_test_containers(actx):
+        if (isinstance(actx, CupyArrayContext)
+              and isinstance(c, (int, float, complex))):
+            # CupyArrayContext does not support zeros_like/ones_like with
+            # Python scalars.
+            continue
         result = getattr(actx.np, sym_name)(c)
         result = actx.thaw(actx.freeze(result))
 
@@ -935,11 +943,13 @@ def _check_allclose(f, arg1, arg2, atol=5.0e-14):
     with pytest.raises(TypeError):
         dc_of_dofs + ary_dof
 
-    bcast_result = ary_dof + bcast_dc_of_dofs
-    bcast_dc_of_dofs + ary_dof
+    if not isinstance(actx, CupyArrayContext):
+        # CupyArrayContext does not support operations between numpy and cupy arrays.
+        bcast_result = ary_dof + bcast_dc_of_dofs
+        bcast_dc_of_dofs + ary_dof
 
-    assert actx.to_numpy(actx.np.linalg.norm(bcast_result.mass
-                                             - 2*ary_of_dofs)) < 1e-8
+        assert actx.to_numpy(actx.np.linalg.norm(bcast_result.mass
+                                                - 2*ary_of_dofs)) < 1e-8
 
     mock_gradient = MyContainerDOFBcast(
             name="yo",
@@ -1129,6 +1139,9 @@ def test_flatten_with_leaf_class(actx_factory):
 
 def test_numpy_conversion(actx_factory):
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("Irrelevant tests for CupyArrayContext. "
+                    "Also, CupyArrayContextdoes not support object arrays.")
 
     nelements = 42
     ac = MyContainer(
@@ -1219,6 +1232,8 @@ def scale_and_orthogonalize(alpha, vel):
 
 def test_actx_compile(actx_factory):
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext does not support object arrays")
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
 
@@ -1236,6 +1251,8 @@ def test_actx_compile(actx_factory):
 
 def test_actx_compile_python_scalar(actx_factory):
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext does not support object arrays")
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
 
@@ -1253,6 +1270,8 @@ def test_actx_compile_python_scalar(actx_factory):
 
 def test_actx_compile_kwargs(actx_factory):
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext does not support object arrays")
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
 
@@ -1273,6 +1292,8 @@ def test_actx_compile_with_tuple_output_keys(actx_factory):
     # key stringification logic.
     from arraycontext import from_numpy, to_numpy
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext does not support object arrays")
 
     def my_rhs(scale, vel):
         result = np.empty((1, 1), dtype=object)
@@ -1337,6 +1358,8 @@ def array_context(self):
 def test_leaf_array_type_broadcasting(actx_factory):
     # test support for https://github.com/inducer/arraycontext/issues/49
     actx = actx_factory()
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext has no leaf array type broadcasting support")
 
     foo = Foo(DOFArray(actx, (actx.zeros(3, dtype=np.float64) + 41, )))
     bar = foo + 4
@@ -1550,6 +1573,9 @@ def test_tagging(actx_factory):
     if isinstance(actx, EagerJAXArrayContext):
         pytest.skip("Eager JAX has no tagging support")
 
+    if isinstance(actx, CupyArrayContext):
+        pytest.skip("CupyArrayContext has no tagging support")
+
     from pytools.tag import Tag
 
     class ExampleTag(Tag):
@@ -1596,6 +1622,9 @@ def test_linspace(actx_factory, args, kwargs):
 
     actx = actx_factory()
 
+    if isinstance(actx, CupyArrayContext) and kwargs.get("dtype") == np.complex128:
+        pytest.skip("CupyArrayContext does not support complex args to linspace")
+
     actx_linspace = actx.to_numpy(actx.np.linspace(*args, **kwargs))
     np_linspace = np.linspace(*args, **kwargs)
 

From 96b7a3d4e6002bf87ee6dacd10c1f758af27f1a9 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 16 Feb 2024 17:36:55 -0600
Subject: [PATCH 02/34] print device name in test

---
 arraycontext/pytest.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py
index f8de488a..6f8b854b 100644
--- a/arraycontext/pytest.py
+++ b/arraycontext/pytest.py
@@ -238,7 +238,10 @@ def __call__(self):
         return CupyArrayContext()
 
     def __str__(self):
-        return "<CupyArrayContext>"
+        import cupy
+        d = cupy.cuda.runtime.getDeviceProperties(cupy.cuda.Device())
+        name = d["name"].decode("utf-8")
+        return f"<CupyArrayContext> on {cupy.cuda.Device()}:{name}"
 
 
 _ARRAY_CONTEXT_FACTORY_REGISTRY: \

From 8dee38d7d7d9176cdf0e5a1da869f4f34e0cf587 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 16 Feb 2024 19:32:03 -0600
Subject: [PATCH 03/34] pylint

---
 arraycontext/pytest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py
index 6f8b854b..d23e1fdb 100644
--- a/arraycontext/pytest.py
+++ b/arraycontext/pytest.py
@@ -238,7 +238,7 @@ def __call__(self):
         return CupyArrayContext()
 
     def __str__(self):
-        import cupy
+        import cupy  # pylint: disable=import-error
         d = cupy.cuda.runtime.getDeviceProperties(cupy.cuda.Device())
         name = d["name"].decode("utf-8")
         return f"<CupyArrayContext> on {cupy.cuda.Device()}:{name}"

From bfa648a5fe6de207f40182a5375f7aefc6553cd2 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 15:40:54 -0500
Subject: [PATCH 04/34] update with current numpy actx

---
 arraycontext/impl/cupy/__init__.py   | 110 +++++++++++++++++----------
 arraycontext/impl/cupy/fake_numpy.py |  85 ++++++++++++---------
 2 files changed, 120 insertions(+), 75 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index f58d16d6..f6c20dea 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -1,11 +1,14 @@
+from __future__ import annotations
+
+
 """
 .. currentmodule:: arraycontext
 
-
 A mod :`cupy`-based array context.
 
 .. autoclass:: CupyArrayContext
 """
+
 __copyright__ = """
 Copyright (C) 2024 University of Illinois Board of Trustees
 """
@@ -30,34 +33,47 @@
 THE SOFTWARE.
 """
 
-from collections.abc import Mapping
+from typing import Any
 
-
-try:
-    import cupy as cp  # type: ignore[import-untyped]
-except ModuleNotFoundError:
-    pass
+import cupy as cp
 
 import loopy as lp
+from pytools.tag import ToTagSetConvertible
+
+from arraycontext.container.traversal import rec_map_array_container, with_array_context
+from arraycontext.context import (
+    Array,
+    ArrayContext,
+    ArrayOrContainerOrScalar,
+    ArrayOrContainerOrScalarT,
+    NumpyOrContainerOrScalar,
+    UntransformedCodeWarning,
+)
 
-from arraycontext.container.traversal import (
-    rec_map_array_container, with_array_context)
-from arraycontext.context import ArrayContext
+
+class CupyNonObjectArrayMetaclass(type):
+    def __instancecheck__(cls, instance: Any) -> bool:
+        return isinstance(instance, cp.ndarray) and instance.dtype != object
+
+
+class CupyNonObjectArray(metaclass=CupyNonObjectArrayMetaclass):
+    pass
 
 
 class CupyArrayContext(ArrayContext):
     """
-    A :class:`ArrayContext` that uses :mod:`cupy.ndarray` to represent arrays
-
+    A :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays.
 
     .. automethod:: __init__
     """
-    def __init__(self):
+
+    _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]
+
+    def __init__(self) -> None:
         super().__init__()
-        self._loopy_transform_cache: \
-                Mapping["lp.TranslationUnit", "lp.TranslationUnit"] = {}
+        self._loopy_transform_cache = {}
 
-        self.array_types = (cp.ndarray,)
+    array_types = (CupyNonObjectArray,)
 
     def _get_fake_numpy_namespace(self):
         from .fake_numpy import CupyFakeNumpyNamespace
@@ -68,29 +84,28 @@ def _get_fake_numpy_namespace(self):
     def clone(self):
         return type(self)()
 
-    def empty(self, shape, dtype):
-        return cp.empty(shape, dtype=dtype)
-
-    def zeros(self, shape, dtype):
-        return cp.zeros(shape, dtype)
+    def from_numpy(self,
+                   array: NumpyOrContainerOrScalar
+                   ) -> ArrayOrContainerOrScalar:
+        return cp.array(array)
 
-    def from_numpy(self, np_array):
-        return cp.array(np_array)
-
-    def to_numpy(self, array):
+    def to_numpy(self,
+                 array: ArrayOrContainerOrScalar
+                 ) -> NumpyOrContainerOrScalar:
         return cp.asnumpy(array)
 
-    def call_loopy(self, t_unit, **kwargs):
+    def call_loopy(
+                self,
+                t_unit: lp.TranslationUnit, **kwargs: Any
+            ) -> dict[str, Array]:
         t_unit = t_unit.copy(target=lp.ExecutableCTarget())
         try:
-            t_unit = self._loopy_transform_cache[t_unit]
+            executor = self._loopy_transform_cache[t_unit]
         except KeyError:
-            orig_t_unit = t_unit
-            t_unit = self.transform_loopy_program(t_unit)
-            self._loopy_transform_cache[orig_t_unit] = t_unit
-            del orig_t_unit
+            executor = self.transform_loopy_program(t_unit).executor()
+            self._loopy_transform_cache[t_unit] = executor
 
-        _, result = t_unit(**kwargs)
+        _, result = executor(**kwargs)
 
         return result
 
@@ -109,15 +124,32 @@ def _thaw(ary):
     # }}}
 
     def transform_loopy_program(self, t_unit):
-        raise ValueError("CupyArrayContext does not implement "
-                         "transform_loopy_program. Sub-classes are supposed "
-                         "to implement it.")
-
-    def tag(self, tags, array):
-        # No tagging support in CupyArrayContext
+        from warnings import warn
+        warn("Using the base "
+                f"{type(self).__name__}.transform_loopy_program "
+                "to transform a translation unit. "
+                "This is a no-op and will result in unoptimized C code for"
+                "the requested optimization, all in a single statement."
+                "This will work, but is unlikely to be performant."
+                f"Instead, subclass {type(self).__name__} and implement "
+                "the specific transform logic required to transform the program "
+                "for your package or application. Check higher-level packages "
+                "(e.g. meshmode), which may already have subclasses you may want "
+                "to build on.",
+                UntransformedCodeWarning, stacklevel=2)
+
+        return t_unit
+
+    def tag(self,
+            tags: ToTagSetConvertible,
+            array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
+        # Cupy (like numpy) doesn't support tagging
         return array
 
-    def tag_axis(self, iaxis, tags, array):
+    def tag_axis(self,
+                 iaxis: int, tags: ToTagSetConvertible,
+                 array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
+        # Cupy (like numpy) doesn't support tagging
         return array
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 3e6fa805..d5cf1ff3 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -23,15 +23,20 @@
 """
 from functools import partial, reduce
 
-import cupy as cp  # type: ignore[import-untyped]  # pylint: disable=import-error
+import cupy as cp
 
-from arraycontext.container import is_array_container
+from arraycontext.container import NotAnArrayContainerError, serialize_container
 from arraycontext.container.traversal import (
-    multimap_reduce_array_container, rec_map_array_container,
-    rec_map_reduce_array_container, rec_multimap_array_container,
-    rec_multimap_reduce_array_container)
+    rec_map_array_container,
+    rec_map_reduce_array_container,
+    rec_multimap_array_container,
+    rec_multimap_reduce_array_container,
+)
+from arraycontext.context import Array, ArrayOrContainer
 from arraycontext.fake_numpy import (
-    BaseFakeNumpyLinalgNamespace, BaseFakeNumpyNamespace)
+    BaseFakeNumpyLinalgNamespace,
+    BaseFakeNumpyNamespace,
+)
 
 
 class CupyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
@@ -39,20 +44,22 @@ class CupyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
     pass
 
 
-_NUMPY_UFUNCS = {"abs", "sin", "cos", "tan", "arcsin", "arccos", "arctan",
-                 "sinh", "cosh", "tanh", "exp", "log", "log10", "isnan",
-                 "sqrt", "concatenate", "transpose",
-                 "ones_like", "maximum", "minimum", "where", "conj", "arctan2",
-                 }
+_NUMPY_UFUNCS = frozenset({"concatenate", "reshape", "transpose",
+                 "ones_like", "where",
+                 *BaseFakeNumpyNamespace._numpy_math_functions
+                 })
 
 
 class CupyFakeNumpyNamespace(BaseFakeNumpyNamespace):
     """
-    A :mod:`numpy` mimic for :class:`CupyArrayContext`.
+    A :mod:`cupy` mimic for :class:`CupyArrayContext`.
     """
     def _get_fake_numpy_linalg_namespace(self):
         return CupyFakeNumpyLinalgNamespace(self._array_context)
 
+    def zeros(self, shape, dtype):
+        return cp.zeros(shape, dtype)
+
     def __getattr__(self, name):
 
         if name in _NUMPY_UFUNCS:
@@ -60,7 +67,7 @@ def __getattr__(self, name):
             return partial(rec_multimap_array_container,
                            getattr(cp, name))
 
-        raise NotImplementedError
+        raise AttributeError(name)
 
     def sum(self, a, axis=None, dtype=None):
         return rec_map_reduce_array_container(sum, partial(cp.sum,
@@ -109,10 +116,7 @@ def less_equal(self, x, y):
     def ravel(self, a, order="C"):
         return rec_map_array_container(partial(cp.ravel, order=order), a)
 
-    def vdot(self, x, y, dtype=None):
-        if dtype is not None:
-            raise NotImplementedError("only 'dtype=None' supported.")
-
+    def vdot(self, x, y):
         return rec_multimap_reduce_array_container(sum, cp.vdot, x, y)
 
     def any(self, a):
@@ -123,34 +127,43 @@ def all(self, a):
         return rec_map_reduce_array_container(partial(reduce, cp.logical_and),
                                               lambda subary: cp.all(subary), a)
 
-    def array_equal(self, a, b):
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
+        false_ary = cp.array(False)
+        true_ary = cp.array(True)
         if type(a) is not type(b):
-            return False
-        elif not is_array_container(a):
-            if a.shape != b.shape:
-                return False
-            else:
-                return cp.all(cp.equal(a, b))
+            return false_ary
+
+        try:
+            serialized_x = serialize_container(a)
+            serialized_y = serialize_container(b)
+        except NotAnArrayContainerError:
+            assert isinstance(a, cp.ndarray)
+            assert isinstance(b, cp.ndarray)
+            return cp.array(cp.array_equal(a, b))
         else:
-            try:
-                return multimap_reduce_array_container(partial(reduce,
-                                                           cp.logical_and),
-                                                   self.array_equal, a, b)
-            except TypeError:
-                return True
+            if len(serialized_x) != len(serialized_y):
+                return false_ary
+            return reduce(
+                    cp.logical_and,
+                    [(true_ary if kx_i == ky_i else false_ary)
+                        and self.array_equal(x_i, y_i)
+                        for (kx_i, x_i), (ky_i, y_i)
+                        in zip(serialized_x, serialized_y)],
+                    true_ary)
+
+    def arange(self, *args, **kwargs):
+        return cp.arange(*args, **kwargs)
+
+    def linspace(self, *args, **kwargs):
+        return cp.linspace(*args, **kwargs)
 
     def zeros_like(self, ary):
-        return rec_multimap_array_container(cp.zeros_like, ary)
+        return rec_map_array_container(cp.zeros_like, ary)
 
     def reshape(self, a, newshape, order="C"):
         return rec_map_array_container(
                 lambda ary: ary.reshape(newshape, order=order),
                 a)
 
-    def arange(self, *args, **kwargs):
-        return cp.arange(*args, **kwargs)
-
-    def linspace(self, *args, **kwargs):
-        return cp.linspace(*args, **kwargs)
 
 # vim: fdm=marker

From 27e5a19bd5dbb788ec4be86f11df8d250e50cf88 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:07:02 -0500
Subject: [PATCH 05/34] restore some tests

---
 arraycontext/impl/cupy/fake_numpy.py | 11 ++++++++++-
 test/test_arraycontext.py            | 26 ++++++--------------------
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index d5cf1ff3..87aa9a21 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -45,7 +45,7 @@ class CupyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
 
 
 _NUMPY_UFUNCS = frozenset({"concatenate", "reshape", "transpose",
-                 "ones_like", "where",
+                 "where",
                  *BaseFakeNumpyNamespace._numpy_math_functions
                  })
 
@@ -158,8 +158,17 @@ def linspace(self, *args, **kwargs):
         return cp.linspace(*args, **kwargs)
 
     def zeros_like(self, ary):
+        if isinstance(ary, (int, float, complex)):
+            # Cupy does not support zeros_like with scalar arguments
+            ary=cp.array(ary)
         return rec_map_array_container(cp.zeros_like, ary)
 
+    def ones_like(self, ary):
+        if isinstance(ary, (int, float, complex)):
+            # Cupy does not support ones_like with scalar arguments
+            ary=cp.array(ary)
+        return rec_map_array_container(cp.ones_like, ary)
+
     def reshape(self, a, newshape, order="C"):
         return rec_map_array_container(
                 lambda ary: ary.reshape(newshape, order=order),
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index 7bda017e..16571372 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -436,12 +436,6 @@ def test_array_context_np_like(actx_factory, sym_name, n_args, dtype):
             actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
 
     for c in (42.0, * _get_test_containers(actx)):
-        if (isinstance(actx, CupyArrayContext)
-              and isinstance(c, (int, float, complex))):
-            # CupyArrayContext does not support zeros_like/ones_like with
-            # Python scalars.
-            continue
-
         result = getattr(actx.np, sym_name)(c)
         result = actx.thaw(actx.freeze(result))
 
@@ -967,13 +961,11 @@ def _check_allclose(f, arg1, arg2, atol=5.0e-14):
     with pytest.raises(TypeError):
         dc_of_dofs + ary_dof
 
-    if not isinstance(actx, CupyArrayContext):
-        # CupyArrayContext does not support operations between numpy and cupy arrays.
-        bcast_result = ary_dof + bcast_dc_of_dofs
-        bcast_dc_of_dofs + ary_dof
+    bcast_result = ary_dof + bcast_dc_of_dofs
+    bcast_dc_of_dofs + ary_dof
 
-        assert actx.to_numpy(actx.np.linalg.norm(bcast_result.mass
-                                                - 2*ary_of_dofs)) < 1e-8
+    assert actx.to_numpy(actx.np.linalg.norm(bcast_result.mass
+                                            - 2*ary_of_dofs)) < 1e-8
 
     mock_gradient = MyContainerDOFBcast(
             name="yo",
@@ -1166,7 +1158,7 @@ def test_numpy_conversion(actx_factory):
     actx = actx_factory()
     if isinstance(actx, CupyArrayContext):
         pytest.skip("Irrelevant tests for CupyArrayContext. "
-                    "Also, CupyArrayContextdoes not support object arrays.")
+                    "Also, CupyArrayContext does not support object arrays.")
     rng = np.random.default_rng()
 
     nelements = 42
@@ -1398,9 +1390,6 @@ def array_context(self):
 def test_no_leaf_array_type_broadcasting(actx_factory):
     # test lack of support for https://github.com/inducer/arraycontext/issues/49
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext has no leaf array type broadcasting support")
-
     dof_ary = DOFArray(actx, (actx.np.zeros(3, dtype=np.float64) + 41, ))
     foo = Foo(dof_ary)
 
@@ -1586,12 +1575,9 @@ def test_to_numpy_on_frozen_arrays(actx_factory):
 def test_tagging(actx_factory):
     actx = actx_factory()
 
-    if isinstance(actx, (NumpyArrayContext, EagerJAXArrayContext)):
+    if isinstance(actx, (NumpyArrayContext, EagerJAXArrayContext, CupyArrayContext)):
         pytest.skip(f"{type(actx)} has no tagging support")
 
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext has no tagging support")
-
     from pytools.tag import Tag
 
     class ExampleTag(Tag):

From 6d507e189838e5657535851852461bc497c71933 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:09:58 -0500
Subject: [PATCH 06/34] ruff

---
 arraycontext/__init__.py             | 2 +-
 arraycontext/impl/cupy/fake_numpy.py | 4 ++--
 arraycontext/pytest.py               | 2 +-
 test/test_arraycontext.py            | 1 -
 4 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/arraycontext/__init__.py b/arraycontext/__init__.py
index 22b6dd0d..fb13b2d1 100644
--- a/arraycontext/__init__.py
+++ b/arraycontext/__init__.py
@@ -105,8 +105,8 @@
     "ArrayOrContainerOrScalarT",
     "ArrayOrContainerT",
     "ArrayT",
-    "CupyArrayContext",
     "CommonSubexpressionTag",
+    "CupyArrayContext",
     "EagerJAXArrayContext",
     "ElementwiseMapKernelTag",
     "NotAnArrayContainerError",
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 87aa9a21..6bf40a00 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -160,13 +160,13 @@ def linspace(self, *args, **kwargs):
     def zeros_like(self, ary):
         if isinstance(ary, (int, float, complex)):
             # Cupy does not support zeros_like with scalar arguments
-            ary=cp.array(ary)
+            ary = cp.array(ary)
         return rec_map_array_container(cp.zeros_like, ary)
 
     def ones_like(self, ary):
         if isinstance(ary, (int, float, complex)):
             # Cupy does not support ones_like with scalar arguments
-            ary=cp.array(ary)
+            ary = cp.array(ary)
         return rec_map_array_container(cp.ones_like, ary)
 
     def reshape(self, a, newshape, order="C"):
diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py
index e4552250..38ce8d2f 100644
--- a/arraycontext/pytest.py
+++ b/arraycontext/pytest.py
@@ -243,7 +243,7 @@ def __str__(self):
         name = d["name"].decode("utf-8")
         return f"<CupyArrayContext> on {cupy.cuda.Device()}:{name}"
 
- 
+
 # {{{ _PytestArrayContextFactory
 
 class _NumpyArrayContextForTests(NumpyArrayContext):
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index 16571372..b462a33c 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -1298,7 +1298,6 @@ def test_actx_compile_kwargs(actx_factory):
 
     rng = np.random.default_rng()
 
-
     compiled_rhs = actx.compile(scale_and_orthogonalize)
 
     v_x = rng.uniform(size=10)

From 8fb4e0b16842d5ba818639055349f860e36c5872 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:12:52 -0500
Subject: [PATCH 07/34] make cupy import optional

---
 arraycontext/impl/cupy/__init__.py   | 10 ++++++++--
 arraycontext/impl/cupy/fake_numpy.py | 24 ++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index f6c20dea..8053ae71 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -35,8 +35,6 @@
 
 from typing import Any
 
-import cupy as cp
-
 import loopy as lp
 from pytools.tag import ToTagSetConvertible
 
@@ -53,6 +51,7 @@
 
 class CupyNonObjectArrayMetaclass(type):
     def __instancecheck__(cls, instance: Any) -> bool:
+        import cupy as cp
         return isinstance(instance, cp.ndarray) and instance.dtype != object
 
 
@@ -87,11 +86,13 @@ def clone(self):
     def from_numpy(self,
                    array: NumpyOrContainerOrScalar
                    ) -> ArrayOrContainerOrScalar:
+        import cupy as cp
         return cp.array(array)
 
     def to_numpy(self,
                  array: ArrayOrContainerOrScalar
                  ) -> NumpyOrContainerOrScalar:
+        import cupy as cp
         return cp.asnumpy(array)
 
     def call_loopy(
@@ -110,12 +111,16 @@ def call_loopy(
         return result
 
     def freeze(self, array):
+        import cupy as cp
+
         def _freeze(ary):
             return cp.asnumpy(ary)
 
         return with_array_context(rec_map_array_container(_freeze, array), actx=None)
 
     def thaw(self, array):
+        import cupy as cp
+
         def _thaw(ary):
             return cp.array(ary)
 
@@ -153,6 +158,7 @@ def tag_axis(self,
         return array
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
+        import cupy as cp
         return cp.einsum(spec, *args)
 
     @property
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 6bf40a00..7776768f 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -23,8 +23,6 @@
 """
 from functools import partial, reduce
 
-import cupy as cp
-
 from arraycontext.container import NotAnArrayContainerError, serialize_container
 from arraycontext.container.traversal import (
     rec_map_array_container,
@@ -58,9 +56,11 @@ def _get_fake_numpy_linalg_namespace(self):
         return CupyFakeNumpyLinalgNamespace(self._array_context)
 
     def zeros(self, shape, dtype):
+        import cupy as cp
         return cp.zeros(shape, dtype)
 
     def __getattr__(self, name):
+        import cupy as cp
 
         if name in _NUMPY_UFUNCS:
             from functools import partial
@@ -70,64 +70,80 @@ def __getattr__(self, name):
         raise AttributeError(name)
 
     def sum(self, a, axis=None, dtype=None):
+        import cupy as cp
         return rec_map_reduce_array_container(sum, partial(cp.sum,
                                                            axis=axis,
                                                            dtype=dtype),
                                               a)
 
     def min(self, a, axis=None):
+        import cupy as cp
         return rec_map_reduce_array_container(
                 partial(reduce, cp.minimum), partial(cp.amin, axis=axis), a)
 
     def max(self, a, axis=None):
+        import cupy as cp
         return rec_map_reduce_array_container(
                 partial(reduce, cp.maximum), partial(cp.amax, axis=axis), a)
 
     def stack(self, arrays, axis=0):
+        import cupy as cp
         return rec_multimap_array_container(
                 lambda *args: cp.stack(args, axis=axis),
                 *arrays)
 
     def broadcast_to(self, array, shape):
+        import cupy as cp
         return rec_map_array_container(partial(cp.broadcast_to, shape=shape), array)
 
     # {{{ relational operators
 
     def equal(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.equal, x, y)
 
     def not_equal(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.not_equal, x, y)
 
     def greater(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.greater, x, y)
 
     def greater_equal(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.greater_equal, x, y)
 
     def less(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.less, x, y)
 
     def less_equal(self, x, y):
+        import cupy as cp
         return rec_multimap_array_container(cp.less_equal, x, y)
 
     # }}}
 
     def ravel(self, a, order="C"):
+        import cupy as cp
         return rec_map_array_container(partial(cp.ravel, order=order), a)
 
     def vdot(self, x, y):
+        import cupy as cp
         return rec_multimap_reduce_array_container(sum, cp.vdot, x, y)
 
     def any(self, a):
+        import cupy as cp
         return rec_map_reduce_array_container(partial(reduce, cp.logical_or),
                                               lambda subary: cp.any(subary), a)
 
     def all(self, a):
+        import cupy as cp
         return rec_map_reduce_array_container(partial(reduce, cp.logical_and),
                                               lambda subary: cp.all(subary), a)
 
     def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
+        import cupy as cp
         false_ary = cp.array(False)
         true_ary = cp.array(True)
         if type(a) is not type(b):
@@ -152,19 +168,23 @@ def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
                     true_ary)
 
     def arange(self, *args, **kwargs):
+        import cupy as cp
         return cp.arange(*args, **kwargs)
 
     def linspace(self, *args, **kwargs):
+        import cupy as cp
         return cp.linspace(*args, **kwargs)
 
     def zeros_like(self, ary):
         if isinstance(ary, (int, float, complex)):
+            import cupy as cp
             # Cupy does not support zeros_like with scalar arguments
             ary = cp.array(ary)
         return rec_map_array_container(cp.zeros_like, ary)
 
     def ones_like(self, ary):
         if isinstance(ary, (int, float, complex)):
+            import cupy as cp
             # Cupy does not support ones_like with scalar arguments
             ary = cp.array(ary)
         return rec_map_array_container(cp.ones_like, ary)

From be70b67ab1092df5c0a77d339755647f6809a194 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:20:31 -0500
Subject: [PATCH 08/34] CI fixes

---
 .github/workflows/ci.yml             | 1 -
 .test-conda-env-py3.yml              | 1 +
 arraycontext/impl/cupy/__init__.py   | 2 +-
 arraycontext/impl/cupy/fake_numpy.py | 2 +-
 4 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 3217f5dc..6e3eab36 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -55,7 +55,6 @@ jobs:
 
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
-                conda install cupy
                 ./run-mypy.sh
 
     pytest3_pocl:
diff --git a/.test-conda-env-py3.yml b/.test-conda-env-py3.yml
index 66daddad..b723a759 100644
--- a/.test-conda-env-py3.yml
+++ b/.test-conda-env-py3.yml
@@ -15,3 +15,4 @@ dependencies:
 - islpy
 - pip
 - jax
+- cupy
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 8053ae71..83eee7d3 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -51,7 +51,7 @@
 
 class CupyNonObjectArrayMetaclass(type):
     def __instancecheck__(cls, instance: Any) -> bool:
-        import cupy as cp
+        import cupy as cp  # type: ignore[import-untyped]
         return isinstance(instance, cp.ndarray) and instance.dtype != object
 
 
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 7776768f..716c5eaa 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -56,7 +56,7 @@ def _get_fake_numpy_linalg_namespace(self):
         return CupyFakeNumpyLinalgNamespace(self._array_context)
 
     def zeros(self, shape, dtype):
-        import cupy as cp
+        import cupy as cp  # type: ignore[import-untyped]
         return cp.zeros(shape, dtype)
 
     def __getattr__(self, name):

From 677419bf5f29883d96b863ec9a238678c55a3d19 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:23:03 -0500
Subject: [PATCH 09/34] remove a few spurious changes

---
 test/test_arraycontext.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index b462a33c..c80e9b8c 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -435,7 +435,7 @@ def test_array_context_np_like(actx_factory, sym_name, n_args, dtype):
     assert_close_to_numpy(
             actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
 
-    for c in (42.0, * _get_test_containers(actx)):
+    for c in (42.0, *_get_test_containers(actx)):
         result = getattr(actx.np, sym_name)(c)
         result = actx.thaw(actx.freeze(result))
 
@@ -965,7 +965,7 @@ def _check_allclose(f, arg1, arg2, atol=5.0e-14):
     bcast_dc_of_dofs + ary_dof
 
     assert actx.to_numpy(actx.np.linalg.norm(bcast_result.mass
-                                            - 2*ary_of_dofs)) < 1e-8
+                                             - 2*ary_of_dofs)) < 1e-8
 
     mock_gradient = MyContainerDOFBcast(
             name="yo",
@@ -1317,7 +1317,6 @@ def test_actx_compile_with_tuple_output_keys(actx_factory):
     # key stringification logic.
     from arraycontext import from_numpy, to_numpy
     actx = actx_factory()
-
     if isinstance(actx, CupyArrayContext):
         pytest.skip("CupyArrayContext does not support object arrays")
 
@@ -1389,6 +1388,7 @@ def array_context(self):
 def test_no_leaf_array_type_broadcasting(actx_factory):
     # test lack of support for https://github.com/inducer/arraycontext/issues/49
     actx = actx_factory()
+
     dof_ary = DOFArray(actx, (actx.np.zeros(3, dtype=np.float64) + 41, ))
     foo = Foo(dof_ary)
 

From 625021168b17c2ecce420f26c22d1fe2b8b23d93 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:27:11 -0500
Subject: [PATCH 10/34] change CI cupy integration

---
 .github/workflows/ci.yml | 5 +++++
 .test-conda-env-py3.yml  | 1 -
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 6e3eab36..2e337e0f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -36,6 +36,10 @@ jobs:
             run: |
                 USE_CONDA_BUILD=1
                 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
+
+                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+                echo "- cupy" >> "$CONDA_ENVIRONMENT"
+
                 . ./prepare-and-run-pylint.sh "$(basename $GITHUB_REPOSITORY)" examples/*.py test/test_*.py
 
     mypy:
@@ -55,6 +59,7 @@ jobs:
 
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
+                conda install cupy
                 ./run-mypy.sh
 
     pytest3_pocl:
diff --git a/.test-conda-env-py3.yml b/.test-conda-env-py3.yml
index b723a759..66daddad 100644
--- a/.test-conda-env-py3.yml
+++ b/.test-conda-env-py3.yml
@@ -15,4 +15,3 @@ dependencies:
 - islpy
 - pip
 - jax
-- cupy

From d61f0cf0784df62b26e9f112dca473fc5db0ad69 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 6 Sep 2024 16:34:35 -0500
Subject: [PATCH 11/34] simplify CI install slightly

---
 .github/workflows/ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 2e337e0f..08d80c1d 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -56,10 +56,10 @@ jobs:
                 # which conflicts with our mypy.python_version = '3.8' setting
                 CONDA_ENVIRONMENT=.test-conda-env-py3.yml
                 sed -i "s/jax/jax<0.4.31/" "$CONDA_ENVIRONMENT"
+                echo "- cupy" >> "$CONDA_ENVIRONMENT"
 
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
-                conda install cupy
                 ./run-mypy.sh
 
     pytest3_pocl:

From fd95813bfe1ac0a303b30ae6700614d309a50ed1 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 3 Dec 2024 16:14:44 -0600
Subject: [PATCH 12/34] lint

---
 .github/workflows/ci.yml             |  3 +++
 arraycontext/impl/cupy/__init__.py   | 21 ++++++++++++++++++++-
 arraycontext/impl/cupy/fake_numpy.py |  9 ++++++---
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f4bbbc59..fe75e8e2 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -56,6 +56,9 @@ jobs:
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
 
+                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+                echo "- cupy" >> "$CONDA_ENVIRONMENT"
+
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
                 ./run-mypy.sh
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 83eee7d3..9556cad3 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -33,7 +33,9 @@
 THE SOFTWARE.
 """
 
-from typing import Any
+from typing import Any, overload
+
+import numpy as np
 
 import loopy as lp
 from pytools.tag import ToTagSetConvertible
@@ -44,6 +46,7 @@
     ArrayContext,
     ArrayOrContainerOrScalar,
     ArrayOrContainerOrScalarT,
+    ContainerOrScalarT,
     NumpyOrContainerOrScalar,
     UntransformedCodeWarning,
 )
@@ -83,12 +86,28 @@ def _get_fake_numpy_namespace(self):
     def clone(self):
         return type(self)()
 
+    @overload
+    def from_numpy(self, array: np.ndarray) -> Array:
+        ...
+
+    @overload
+    def from_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT:
+        ...
+
     def from_numpy(self,
                    array: NumpyOrContainerOrScalar
                    ) -> ArrayOrContainerOrScalar:
         import cupy as cp
         return cp.array(array)
 
+    @overload
+    def to_numpy(self, array: Array) -> np.ndarray:
+        ...
+
+    @overload
+    def to_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT:
+        ...
+
     def to_numpy(self,
                  array: ArrayOrContainerOrScalar
                  ) -> NumpyOrContainerOrScalar:
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 716c5eaa..5d9bcade 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+
 __copyright__ = """
 Copyright (C) 2024 University of Illinois Board of Trustees
 """
@@ -164,7 +167,7 @@ def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
                     [(true_ary if kx_i == ky_i else false_ary)
                         and self.array_equal(x_i, y_i)
                         for (kx_i, x_i), (ky_i, y_i)
-                        in zip(serialized_x, serialized_y)],
+                        in zip(serialized_x, serialized_y, strict=True)],
                     true_ary)
 
     def arange(self, *args, **kwargs):
@@ -176,14 +179,14 @@ def linspace(self, *args, **kwargs):
         return cp.linspace(*args, **kwargs)
 
     def zeros_like(self, ary):
-        if isinstance(ary, (int, float, complex)):
+        if isinstance(ary, int | float | complex):
             import cupy as cp
             # Cupy does not support zeros_like with scalar arguments
             ary = cp.array(ary)
         return rec_map_array_container(cp.zeros_like, ary)
 
     def ones_like(self, ary):
-        if isinstance(ary, (int, float, complex)):
+        if isinstance(ary, int | float | complex):
             import cupy as cp
             # Cupy does not support ones_like with scalar arguments
             ary = cp.array(ary)

From 5f4c4d9e5befa8f713e8e670179e09d7d614110a Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 3 Dec 2024 16:42:22 -0600
Subject: [PATCH 13/34] update docs

---
 README.rst                             |  1 +
 arraycontext/impl/cupy/__init__.py     | 10 +++++-----
 arraycontext/impl/numpy/__init__.py    |  2 +-
 arraycontext/impl/pyopencl/__init__.py |  2 +-
 doc/Makefile                           |  2 +-
 doc/conf.py                            |  1 +
 doc/implementations.rst                |  6 ++++++
 doc/index.rst                          |  3 ++-
 8 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/README.rst b/README.rst
index a704c122..1206093e 100644
--- a/README.rst
+++ b/README.rst
@@ -16,6 +16,7 @@ code to work with all of them? No problem! Comes with pre-made array context
 implementations for:
 
 - numpy
+- cupy
 - `PyOpenCL <https://documen.tician.de/pyopencl/array.html>`__
 - `JAX <https://jax.readthedocs.io/en/latest/>`__
 - `Pytato <https://documen.tician.de/pytato>`__ (for lazy/deferred evaluation)
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 9556cad3..87d41570 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -1,14 +1,14 @@
-from __future__ import annotations
-
-
 """
 .. currentmodule:: arraycontext
 
-A mod :`cupy`-based array context.
+A :mod:`cupy`-based array context.
 
 .. autoclass:: CupyArrayContext
 """
 
+from __future__ import annotations
+
+
 __copyright__ = """
 Copyright (C) 2024 University of Illinois Board of Trustees
 """
@@ -64,7 +64,7 @@ class CupyNonObjectArray(metaclass=CupyNonObjectArrayMetaclass):
 
 class CupyArrayContext(ArrayContext):
     """
-    A :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays.
+    An :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays.
 
     .. automethod:: __init__
     """
diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py
index f9d6c541..6fcc88c7 100644
--- a/arraycontext/impl/numpy/__init__.py
+++ b/arraycontext/impl/numpy/__init__.py
@@ -63,7 +63,7 @@ class NumpyNonObjectArray(metaclass=NumpyNonObjectArrayMetaclass):
 
 class NumpyArrayContext(ArrayContext):
     """
-    A :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays.
+    An :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays.
 
     .. automethod:: __init__
     """
diff --git a/arraycontext/impl/pyopencl/__init__.py b/arraycontext/impl/pyopencl/__init__.py
index 84d5f483..190599e2 100644
--- a/arraycontext/impl/pyopencl/__init__.py
+++ b/arraycontext/impl/pyopencl/__init__.py
@@ -58,7 +58,7 @@
 
 class PyOpenCLArrayContext(ArrayContext):
     """
-    A :class:`ArrayContext` that uses :class:`pyopencl.array.Array` instances
+    An :class:`ArrayContext` that uses :class:`pyopencl.array.Array` instances
     for its base array class.
 
     .. attribute:: context
diff --git a/doc/Makefile b/doc/Makefile
index d0ac5f2f..0568a00c 100644
--- a/doc/Makefile
+++ b/doc/Makefile
@@ -3,7 +3,7 @@
 
 # You can set these variables from the command line, and also
 # from the environment for the first two.
-SPHINXOPTS    ?=
+SPHINXOPTS    ?= -W -n
 SPHINXBUILD   ?= python $(shell which sphinx-build)
 SOURCEDIR     = .
 BUILDDIR      = _build
diff --git a/doc/conf.py b/doc/conf.py
index 0042ae57..04b3971b 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -23,6 +23,7 @@
     "pytest": ("https://docs.pytest.org/en/latest/", None),
     "python": ("https://docs.python.org/3/", None),
     "pytools": ("https://documen.tician.de/pytools", None),
+    "cupy": ("https://docs.cupy.dev/en/stable/", None),
 }
 
 # Some modules need to import things just so that sphinx can resolve symbols in
diff --git a/doc/implementations.rst b/doc/implementations.rst
index 4023e37c..2e6344e8 100644
--- a/doc/implementations.rst
+++ b/doc/implementations.rst
@@ -13,6 +13,12 @@ Array context based on :mod:`numpy`
 
 .. automodule:: arraycontext.impl.numpy
 
+
+Array context based on :mod:`cupy`
+--------------------------------------------
+
+.. automodule:: arraycontext.impl.cupy
+
 Array context based on :mod:`pyopencl.array`
 --------------------------------------------
 
diff --git a/doc/index.rst b/doc/index.rst
index d3f9854b..2445d2e2 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -6,6 +6,7 @@ code to work with all of them? No problem! Comes with pre-made array context
 implementations for:
 
 - :mod:`numpy`
+- :mod:`cupy`
 - :mod:`pyopencl`
 - :mod:`jax.numpy`
 - :mod:`pytato` (for lazy/deferred evaluation)
@@ -13,7 +14,7 @@ implementations for:
 - Profiling
 
 :mod:`arraycontext` started life as an array abstraction for use with the
-:mod:`meshmode` unstrucuted discretization package.
+:mod:`meshmode` unstructured discretization package.
 
 Design Guidelines
 -----------------

From 2296c6d93d5f9439c3fa593401bd58360021654e Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 31 Jan 2025 16:46:21 -0600
Subject: [PATCH 14/34] improve array container support in {to,from}_numpy

---
 arraycontext/impl/cupy/__init__.py | 14 ++++++++++++--
 test/test_arraycontext.py          | 16 +---------------
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 87d41570..27888abb 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -98,7 +98,12 @@ def from_numpy(self,
                    array: NumpyOrContainerOrScalar
                    ) -> ArrayOrContainerOrScalar:
         import cupy as cp
-        return cp.array(array)
+
+        def _from_numpy(ary):
+            return cp.array(ary)
+
+        return with_array_context(rec_map_array_container(_from_numpy, array),
+                                  actx=self)
 
     @overload
     def to_numpy(self, array: Array) -> np.ndarray:
@@ -112,7 +117,12 @@ def to_numpy(self,
                  array: ArrayOrContainerOrScalar
                  ) -> NumpyOrContainerOrScalar:
         import cupy as cp
-        return cp.asnumpy(array)
+
+        def _to_numpy(ary):
+            return cp.asnumpy(ary)
+
+        return with_array_context(rec_map_array_container(_to_numpy, array),
+                                  actx=self)
 
     def call_loopy(
                 self,
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index 9324e18a..1307acb0 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -1014,9 +1014,6 @@ def test_flatten_with_leaf_class(actx_factory):
 
 def test_numpy_conversion(actx_factory):
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("Irrelevant tests for CupyArrayContext. "
-                    "Also, CupyArrayContext does not support object arrays.")
     rng = np.random.default_rng()
 
     nelements = 42
@@ -1033,7 +1030,7 @@ def test_numpy_conversion(actx_factory):
     assert np.allclose(ac.mass, ac_roundtrip.mass)
     assert np.allclose(ac.momentum[0], ac_roundtrip.momentum[0])
 
-    if not isinstance(actx, NumpyArrayContext):
+    if not isinstance(actx, NumpyArrayContext | CupyArrayContext):
         from dataclasses import replace
         ac_with_cl = replace(ac, enthalpy=ac_actx.mass)
         with pytest.raises(TypeError):
@@ -1094,9 +1091,6 @@ def scale_and_orthogonalize(alpha, vel):
 
 def test_actx_compile(actx_factory):
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext does not support object arrays")
-
     rng = np.random.default_rng()
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
@@ -1115,9 +1109,6 @@ def test_actx_compile(actx_factory):
 
 def test_actx_compile_python_scalar(actx_factory):
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext does not support object arrays")
-
     rng = np.random.default_rng()
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
@@ -1136,9 +1127,6 @@ def test_actx_compile_python_scalar(actx_factory):
 
 def test_actx_compile_kwargs(actx_factory):
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext does not support object arrays")
-
     rng = np.random.default_rng()
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
@@ -1160,8 +1148,6 @@ def test_actx_compile_with_tuple_output_keys(actx_factory):
     # key stringification logic.
     from arraycontext import from_numpy, to_numpy
     actx = actx_factory()
-    if isinstance(actx, CupyArrayContext):
-        pytest.skip("CupyArrayContext does not support object arrays")
 
     rng = np.random.default_rng()
 

From 8fd5488b637c6a19da9a5ba0152157d46e253a45 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 31 Jan 2025 16:51:28 -0600
Subject: [PATCH 15/34] WS fix

---
 test/test_arraycontext.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
index 1307acb0..7d7e5a21 100644
--- a/test/test_arraycontext.py
+++ b/test/test_arraycontext.py
@@ -1148,7 +1148,6 @@ def test_actx_compile_with_tuple_output_keys(actx_factory):
     # key stringification logic.
     from arraycontext import from_numpy, to_numpy
     actx = actx_factory()
-
     rng = np.random.default_rng()
 
     def my_rhs(scale, vel):

From 6f3cd949405247a5511c85add95a81fe437f8d40 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 3 Feb 2025 15:06:10 -0600
Subject: [PATCH 16/34] fixes

---
 arraycontext/impl/cupy/__init__.py   | 2 +-
 arraycontext/impl/cupy/fake_numpy.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 27888abb..610c6ff5 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -122,7 +122,7 @@ def _to_numpy(ary):
             return cp.asnumpy(ary)
 
         return with_array_context(rec_map_array_container(_to_numpy, array),
-                                  actx=self)
+                                  actx=None)
 
     def call_loopy(
                 self,
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index 5d9bcade..e45ca8a8 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -179,15 +179,15 @@ def linspace(self, *args, **kwargs):
         return cp.linspace(*args, **kwargs)
 
     def zeros_like(self, ary):
+        import cupy as cp
         if isinstance(ary, int | float | complex):
-            import cupy as cp
             # Cupy does not support zeros_like with scalar arguments
             ary = cp.array(ary)
         return rec_map_array_container(cp.zeros_like, ary)
 
     def ones_like(self, ary):
+        import cupy as cp
         if isinstance(ary, int | float | complex):
-            import cupy as cp
             # Cupy does not support ones_like with scalar arguments
             ary = cp.array(ary)
         return rec_map_array_container(cp.ones_like, ary)

From ab8266d971bbb40467bb647d4fe504d67f72e513 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 11:26:29 -0600
Subject: [PATCH 17/34] allow optional device selection

---
 arraycontext/impl/cupy/__init__.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 610c6ff5..c8c8e91a 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -71,10 +71,14 @@ class CupyArrayContext(ArrayContext):
 
     _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]
 
-    def __init__(self) -> None:
+    def __init__(self, device: int | None = None) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
+        if device is not None:
+            import cupy as cp
+            cp.cuda.runtime.setDevice(device)
+
     array_types = (CupyNonObjectArray,)
 
     def _get_fake_numpy_namespace(self):

From 79b0bc37e75d4f58fb9c6f98323fb9a3b158476a Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 11:31:44 -0600
Subject: [PATCH 18/34] try running cupy via gitlab

---
 .gitlab-ci.yml | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 313b7633..94413889 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -33,6 +33,24 @@ Python 3 Nvidia Titan V:
     reports:
       junit: test/pytest.xml
 
+Python 3 CuPy Nvidia Titan V:
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+    echo "- cupy" >> "$CONDA_ENVIRONMENT"
+    build_py_project_in_venv
+    test_py_project
+
+  tags:
+  - python3
+  - nvidia-titan-v
+  except:
+  - tags
+  artifacts:
+    reports:
+      junit: test/pytest.xml
+
 Python 3 POCL Nvidia Titan V:
   script: |
     curl -L -O https://tiker.net/ci-support-v0

From 70aff99b21463f30ab077a4597c02d24a8c1bbe8 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 12:11:53 -0600
Subject: [PATCH 19/34] debug more

---
 .gitlab-ci.yml | 274 ++++++++++++++++++++++++-------------------------
 1 file changed, 137 insertions(+), 137 deletions(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 94413889..cdd06e06 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -1,37 +1,37 @@
-Python 3 POCL:
-  script: |
-    export PYOPENCL_TEST=portable:cpu
-    export EXTRA_INSTALL="jax[cpu]"
-    export JAX_PLATFORMS=cpu
-    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
-    . ./build-and-test-py-project.sh
-  tags:
-  - python3
-  - pocl
-  - mpi
-  except:
-  - tags
-  artifacts:
-    reports:
-      junit: test/pytest.xml
-
-Python 3 Nvidia Titan V:
-  script: |
-    curl -L -O https://tiker.net/ci-support-v0
-    . ./ci-support-v0
-    export PYOPENCL_TEST=nvi:titan
-    build_py_project_in_venv
-    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
-    test_py_project
-
-  tags:
-  - python3
-  - nvidia-titan-v
-  except:
-  - tags
-  artifacts:
-    reports:
-      junit: test/pytest.xml
+# Python 3 POCL:
+#   script: |
+#     export PYOPENCL_TEST=portable:cpu
+#     export EXTRA_INSTALL="jax[cpu]"
+#     export JAX_PLATFORMS=cpu
+#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
+#     . ./build-and-test-py-project.sh
+#   tags:
+#   - python3
+#   - pocl
+#   - mpi
+#   except:
+#   - tags
+#   artifacts:
+#     reports:
+#       junit: test/pytest.xml
+
+# Python 3 Nvidia Titan V:
+#   script: |
+#     curl -L -O https://tiker.net/ci-support-v0
+#     . ./ci-support-v0
+#     export PYOPENCL_TEST=nvi:titan
+#     build_py_project_in_venv
+#     pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
+#     test_py_project
+
+#   tags:
+#   - python3
+#   - nvidia-titan-v
+#   except:
+#   - tags
+#   artifacts:
+#     reports:
+#       junit: test/pytest.xml
 
 Python 3 CuPy Nvidia Titan V:
   script: |
@@ -39,7 +39,7 @@ Python 3 CuPy Nvidia Titan V:
     . ./ci-support-v0
     CONDA_ENVIRONMENT=.test-conda-env-py3.yml
     echo "- cupy" >> "$CONDA_ENVIRONMENT"
-    build_py_project_in_venv
+    build_py_project_in_conda_env
     test_py_project
 
   tags:
@@ -51,105 +51,105 @@ Python 3 CuPy Nvidia Titan V:
     reports:
       junit: test/pytest.xml
 
-Python 3 POCL Nvidia Titan V:
-  script: |
-    curl -L -O https://tiker.net/ci-support-v0
-    . ./ci-support-v0
-    export PYOPENCL_TEST=port:titan
-    build_py_project_in_venv
-    test_py_project
-
-  tags:
-  - python3
-  - nvidia-titan-v
-  except:
-  - tags
-  artifacts:
-    reports:
-      junit: test/pytest.xml
-
-Python 3 POCL Examples:
-  script:
-  - test -n "$SKIP_EXAMPLES" && exit
-  - export PYOPENCL_TEST=portable:cpu
-  - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
-  - ". ./build-py-project-and-run-examples.sh"
-  tags:
-  - python3
-  - pocl
-  - large-node
-  except:
-  - tags
-
-Python 3 Conda:
-  script: |
-    export PYOPENCL_TEST=portable:cpu
-
-    # Avoid crashes like https://gitlab.tiker.net/inducer/arraycontext/-/jobs/536021
-    sed -i 's/jax/jax !=0.4.6/' .test-conda-env-py3.yml
-
-    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
-    . ./build-and-test-py-project-within-miniconda.sh
-  tags:
-  # - docker-runner
-  - linux
-  - large-node
-  except:
-  - tags
-
-Documentation:
-  script: |
-    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-docs.sh
-    CI_SUPPORT_SPHINX_VERSION_SPECIFIER=">=4.0"
-    . ./build-docs.sh
-  tags:
-  - python3
-
-Ruff:
-  script:
-  - pipx install ruff
-  - ruff check
-  tags:
-  - docker-runner
-  except:
-  - tags
-
-Pylint:
-  script: |
-    EXTRA_INSTALL="jax[cpu]"
-    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
-    . ./prepare-and-run-pylint.sh "$CI_PROJECT_NAME" examples/*.py test/test_*.py
-  tags:
-  - python3
-  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:
-    - DOWNSTREAM_PROJECT: [meshmode, grudge, mirgecom, mirgecom_examples]
-  tags:
-  - large-node
-  - "docker-runner"
-  script: |
-    curl -L -O https://tiker.net/ci-support-v0
-    . ./ci-support-v0
-    test_downstream "$DOWNSTREAM_PROJECT"
-
-    if [[ "$DOWNSTREAM_PROJECT" = "meshmode" ]]; then
-        python ../examples/simple-dg.py --lazy
-    fi
+# Python 3 POCL Nvidia Titan V:
+#   script: |
+#     curl -L -O https://tiker.net/ci-support-v0
+#     . ./ci-support-v0
+#     export PYOPENCL_TEST=port:titan
+#     build_py_project_in_venv
+#     test_py_project
+
+#   tags:
+#   - python3
+#   - nvidia-titan-v
+#   except:
+#   - tags
+#   artifacts:
+#     reports:
+#       junit: test/pytest.xml
+
+# Python 3 POCL Examples:
+#   script:
+#   - test -n "$SKIP_EXAMPLES" && exit
+#   - export PYOPENCL_TEST=portable:cpu
+#   - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
+#   - ". ./build-py-project-and-run-examples.sh"
+#   tags:
+#   - python3
+#   - pocl
+#   - large-node
+#   except:
+#   - tags
+
+# Python 3 Conda:
+#   script: |
+#     export PYOPENCL_TEST=portable:cpu
+
+#     # Avoid crashes like https://gitlab.tiker.net/inducer/arraycontext/-/jobs/536021
+#     sed -i 's/jax/jax !=0.4.6/' .test-conda-env-py3.yml
+
+#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
+#     . ./build-and-test-py-project-within-miniconda.sh
+#   tags:
+#   # - docker-runner
+#   - linux
+#   - large-node
+#   except:
+#   - tags
+
+# Documentation:
+#   script: |
+#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-docs.sh
+#     CI_SUPPORT_SPHINX_VERSION_SPECIFIER=">=4.0"
+#     . ./build-docs.sh
+#   tags:
+#   - python3
+
+# Ruff:
+#   script:
+#   - pipx install ruff
+#   - ruff check
+#   tags:
+#   - docker-runner
+#   except:
+#   - tags
+
+# Pylint:
+#   script: |
+#     EXTRA_INSTALL="jax[cpu]"
+#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
+#     . ./prepare-and-run-pylint.sh "$CI_PROJECT_NAME" examples/*.py test/test_*.py
+#   tags:
+#   - python3
+#   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:
+#     - DOWNSTREAM_PROJECT: [meshmode, grudge, mirgecom, mirgecom_examples]
+#   tags:
+#   - large-node
+#   - "docker-runner"
+#   script: |
+#     curl -L -O https://tiker.net/ci-support-v0
+#     . ./ci-support-v0
+#     test_downstream "$DOWNSTREAM_PROJECT"
+
+#     if [[ "$DOWNSTREAM_PROJECT" = "meshmode" ]]; then
+#         python ../examples/simple-dg.py --lazy
+#     fi

From 8b5c6cff62ae0af1b21627d2d2746c92572906a6 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 16:29:07 -0600
Subject: [PATCH 20/34] Revert "debug more"

This reverts commit 70aff99b21463f30ab077a4597c02d24a8c1bbe8.
---
 .gitlab-ci.yml | 274 ++++++++++++++++++++++++-------------------------
 1 file changed, 137 insertions(+), 137 deletions(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index cdd06e06..94413889 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -1,37 +1,37 @@
-# Python 3 POCL:
-#   script: |
-#     export PYOPENCL_TEST=portable:cpu
-#     export EXTRA_INSTALL="jax[cpu]"
-#     export JAX_PLATFORMS=cpu
-#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
-#     . ./build-and-test-py-project.sh
-#   tags:
-#   - python3
-#   - pocl
-#   - mpi
-#   except:
-#   - tags
-#   artifacts:
-#     reports:
-#       junit: test/pytest.xml
-
-# Python 3 Nvidia Titan V:
-#   script: |
-#     curl -L -O https://tiker.net/ci-support-v0
-#     . ./ci-support-v0
-#     export PYOPENCL_TEST=nvi:titan
-#     build_py_project_in_venv
-#     pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
-#     test_py_project
-
-#   tags:
-#   - python3
-#   - nvidia-titan-v
-#   except:
-#   - tags
-#   artifacts:
-#     reports:
-#       junit: test/pytest.xml
+Python 3 POCL:
+  script: |
+    export PYOPENCL_TEST=portable:cpu
+    export EXTRA_INSTALL="jax[cpu]"
+    export JAX_PLATFORMS=cpu
+    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
+    . ./build-and-test-py-project.sh
+  tags:
+  - python3
+  - pocl
+  - mpi
+  except:
+  - tags
+  artifacts:
+    reports:
+      junit: test/pytest.xml
+
+Python 3 Nvidia Titan V:
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    export PYOPENCL_TEST=nvi:titan
+    build_py_project_in_venv
+    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
+    test_py_project
+
+  tags:
+  - python3
+  - nvidia-titan-v
+  except:
+  - tags
+  artifacts:
+    reports:
+      junit: test/pytest.xml
 
 Python 3 CuPy Nvidia Titan V:
   script: |
@@ -39,7 +39,7 @@ Python 3 CuPy Nvidia Titan V:
     . ./ci-support-v0
     CONDA_ENVIRONMENT=.test-conda-env-py3.yml
     echo "- cupy" >> "$CONDA_ENVIRONMENT"
-    build_py_project_in_conda_env
+    build_py_project_in_venv
     test_py_project
 
   tags:
@@ -51,105 +51,105 @@ Python 3 CuPy Nvidia Titan V:
     reports:
       junit: test/pytest.xml
 
-# Python 3 POCL Nvidia Titan V:
-#   script: |
-#     curl -L -O https://tiker.net/ci-support-v0
-#     . ./ci-support-v0
-#     export PYOPENCL_TEST=port:titan
-#     build_py_project_in_venv
-#     test_py_project
-
-#   tags:
-#   - python3
-#   - nvidia-titan-v
-#   except:
-#   - tags
-#   artifacts:
-#     reports:
-#       junit: test/pytest.xml
-
-# Python 3 POCL Examples:
-#   script:
-#   - test -n "$SKIP_EXAMPLES" && exit
-#   - export PYOPENCL_TEST=portable:cpu
-#   - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
-#   - ". ./build-py-project-and-run-examples.sh"
-#   tags:
-#   - python3
-#   - pocl
-#   - large-node
-#   except:
-#   - tags
-
-# Python 3 Conda:
-#   script: |
-#     export PYOPENCL_TEST=portable:cpu
-
-#     # Avoid crashes like https://gitlab.tiker.net/inducer/arraycontext/-/jobs/536021
-#     sed -i 's/jax/jax !=0.4.6/' .test-conda-env-py3.yml
-
-#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
-#     . ./build-and-test-py-project-within-miniconda.sh
-#   tags:
-#   # - docker-runner
-#   - linux
-#   - large-node
-#   except:
-#   - tags
-
-# Documentation:
-#   script: |
-#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-docs.sh
-#     CI_SUPPORT_SPHINX_VERSION_SPECIFIER=">=4.0"
-#     . ./build-docs.sh
-#   tags:
-#   - python3
-
-# Ruff:
-#   script:
-#   - pipx install ruff
-#   - ruff check
-#   tags:
-#   - docker-runner
-#   except:
-#   - tags
-
-# Pylint:
-#   script: |
-#     EXTRA_INSTALL="jax[cpu]"
-#     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
-#     . ./prepare-and-run-pylint.sh "$CI_PROJECT_NAME" examples/*.py test/test_*.py
-#   tags:
-#   - python3
-#   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:
-#     - DOWNSTREAM_PROJECT: [meshmode, grudge, mirgecom, mirgecom_examples]
-#   tags:
-#   - large-node
-#   - "docker-runner"
-#   script: |
-#     curl -L -O https://tiker.net/ci-support-v0
-#     . ./ci-support-v0
-#     test_downstream "$DOWNSTREAM_PROJECT"
-
-#     if [[ "$DOWNSTREAM_PROJECT" = "meshmode" ]]; then
-#         python ../examples/simple-dg.py --lazy
-#     fi
+Python 3 POCL Nvidia Titan V:
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    export PYOPENCL_TEST=port:titan
+    build_py_project_in_venv
+    test_py_project
+
+  tags:
+  - python3
+  - nvidia-titan-v
+  except:
+  - tags
+  artifacts:
+    reports:
+      junit: test/pytest.xml
+
+Python 3 POCL Examples:
+  script:
+  - test -n "$SKIP_EXAMPLES" && exit
+  - export PYOPENCL_TEST=portable:cpu
+  - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
+  - ". ./build-py-project-and-run-examples.sh"
+  tags:
+  - python3
+  - pocl
+  - large-node
+  except:
+  - tags
+
+Python 3 Conda:
+  script: |
+    export PYOPENCL_TEST=portable:cpu
+
+    # Avoid crashes like https://gitlab.tiker.net/inducer/arraycontext/-/jobs/536021
+    sed -i 's/jax/jax !=0.4.6/' .test-conda-env-py3.yml
+
+    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh
+    . ./build-and-test-py-project-within-miniconda.sh
+  tags:
+  # - docker-runner
+  - linux
+  - large-node
+  except:
+  - tags
+
+Documentation:
+  script: |
+    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-docs.sh
+    CI_SUPPORT_SPHINX_VERSION_SPECIFIER=">=4.0"
+    . ./build-docs.sh
+  tags:
+  - python3
+
+Ruff:
+  script:
+  - pipx install ruff
+  - ruff check
+  tags:
+  - docker-runner
+  except:
+  - tags
+
+Pylint:
+  script: |
+    EXTRA_INSTALL="jax[cpu]"
+    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
+    . ./prepare-and-run-pylint.sh "$CI_PROJECT_NAME" examples/*.py test/test_*.py
+  tags:
+  - python3
+  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:
+    - DOWNSTREAM_PROJECT: [meshmode, grudge, mirgecom, mirgecom_examples]
+  tags:
+  - large-node
+  - "docker-runner"
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    test_downstream "$DOWNSTREAM_PROJECT"
+
+    if [[ "$DOWNSTREAM_PROJECT" = "meshmode" ]]; then
+        python ../examples/simple-dg.py --lazy
+    fi

From 8561b2f937112b55feadb1da639a51716d175179 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 16:30:14 -0600
Subject: [PATCH 21/34] only test pocl-cpu in cupy test

---
 .gitlab-ci.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 94413889..2bca9140 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -39,6 +39,7 @@ Python 3 CuPy Nvidia Titan V:
     . ./ci-support-v0
     CONDA_ENVIRONMENT=.test-conda-env-py3.yml
     echo "- cupy" >> "$CONDA_ENVIRONMENT"
+    export PYOPENCL_TEST=port:cpu
     build_py_project_in_venv
     test_py_project
 

From 904e061d07fab44208fd768e92540363d5e270d0 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Fri, 7 Feb 2025 16:38:38 -0600
Subject: [PATCH 22/34] fix conda build

---
 .gitlab-ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 2bca9140..e2fcd3cf 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -40,7 +40,7 @@ Python 3 CuPy Nvidia Titan V:
     CONDA_ENVIRONMENT=.test-conda-env-py3.yml
     echo "- cupy" >> "$CONDA_ENVIRONMENT"
     export PYOPENCL_TEST=port:cpu
-    build_py_project_in_venv
+    build_py_project_in_conda_env
     test_py_project
 
   tags:

From 340f9dc5fbc46cf7f7fafe5608c6380dfa96ec3a Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Sat, 8 Feb 2025 09:49:39 -0600
Subject: [PATCH 23/34] add to coverage table

---
 doc/make_numpy_coverage_table.py |   4 +-
 doc/numpy_coverage.rst           | 108 ++++++++++++++++++++++++++++++-
 2 files changed, 108 insertions(+), 4 deletions(-)

diff --git a/doc/make_numpy_coverage_table.py b/doc/make_numpy_coverage_table.py
index 57f833d7..abaab57c 100644
--- a/doc/make_numpy_coverage_table.py
+++ b/doc/make_numpy_coverage_table.py
@@ -11,7 +11,7 @@
 
 .. code::
 
-    python make_numpy_support_table.py numpy_coverage.rst
+    python make_numpy_coverage_table.py numpy_coverage.rst
 """
 from __future__ import annotations
 
@@ -67,6 +67,8 @@ def initialize_contexts():
             arraycontext.EagerJAXArrayContext(),
             arraycontext.PytatoPyOpenCLArrayContext(queue),
             arraycontext.PytatoJAXArrayContext(),
+            arraycontext.NumpyArrayContext(),
+            arraycontext.CupyArrayContext(),
             ]
 
 
diff --git a/doc/numpy_coverage.rst b/doc/numpy_coverage.rst
index 5a7b2918..679bceb1 100644
--- a/doc/numpy_coverage.rst
+++ b/doc/numpy_coverage.rst
@@ -18,12 +18,18 @@ Array creation routines
       - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
       - :class:`~arraycontext.PytatoJAXArrayContext`
+      - :class:`~arraycontext.NumpyArrayContext`
+      - :class:`~arraycontext.CupyArrayContext`
     * - :func:`numpy.empty_like`
       - :green:`Yes`
       - :green:`Yes`
+      - :red:`No`
+      - :red:`No`
+      - :red:`No`
+      - :red:`No`
+    * - :func:`numpy.ones_like`
       - :green:`Yes`
       - :green:`Yes`
-    * - :func:`numpy.ones_like`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
@@ -33,16 +39,22 @@ Array creation routines
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.full_like`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :red:`No`
+      - :red:`No`
     * - :func:`numpy.copy`
       - :green:`Yes`
       - :green:`Yes`
       - :red:`No`
       - :red:`No`
+      - :red:`No`
+      - :red:`No`
 
 Array manipulation routines
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -55,36 +67,50 @@ Array manipulation routines
       - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
       - :class:`~arraycontext.PytatoJAXArrayContext`
+      - :class:`~arraycontext.NumpyArrayContext`
+      - :class:`~arraycontext.CupyArrayContext`
     * - :func:`numpy.reshape`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.ravel`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.transpose`
       - :red:`No`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.broadcast_to`
       - :red:`No`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.concatenate`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.stack`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
 
 Linear algebra
 ~~~~~~~~~~~~~~
@@ -97,11 +123,15 @@ Linear algebra
       - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
       - :class:`~arraycontext.PytatoJAXArrayContext`
+      - :class:`~arraycontext.NumpyArrayContext`
+      - :class:`~arraycontext.CupyArrayContext`
     * - :func:`numpy.vdot`
       - :green:`Yes`
       - :green:`Yes`
-      - :red:`No`
-      - :red:`No`
+      - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
 
 Logic Functions
 ~~~~~~~~~~~~~~~
@@ -114,46 +144,64 @@ Logic Functions
       - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
       - :class:`~arraycontext.PytatoJAXArrayContext`
+      - :class:`~arraycontext.NumpyArrayContext`
+      - :class:`~arraycontext.CupyArrayContext`
     * - :func:`numpy.all`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.any`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.greater`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.greater_equal`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.less`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.less_equal`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.equal`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.not_equal`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
 
 Mathematical functions
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -166,133 +214,187 @@ Mathematical functions
       - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
       - :class:`~arraycontext.PytatoJAXArrayContext`
+      - :class:`~arraycontext.NumpyArrayContext`
+      - :class:`~arraycontext.CupyArrayContext`
     * - :data:`numpy.sin`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.cos`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.tan`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arcsin`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arccos`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arctan`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arctan2`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.sinh`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.cosh`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.tanh`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.floor`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.ceil`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.sum`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.exp`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.log`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.log10`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.real`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.imag`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.conjugate`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.maximum`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.amax`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :red:`No`
+      - :red:`No`
     * - :data:`numpy.minimum`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.amin`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :red:`No`
+      - :red:`No`
     * - :data:`numpy.sqrt`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.absolute`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.fabs`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`

From 3ee28cc4d998a3379f3d601dafe744abeef2aa13 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 11 Feb 2025 12:52:31 -0800
Subject: [PATCH 24/34] keep device in clone()

Co-authored-by: Alexandru Fikl <alexfikl@gmail.com>
---
 arraycontext/impl/cupy/__init__.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index c8c8e91a..17414d1b 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -75,6 +75,8 @@ def __init__(self, device: int | None = None) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
+        self.device = device
+
         if device is not None:
             import cupy as cp
             cp.cuda.runtime.setDevice(device)
@@ -88,7 +90,7 @@ def _get_fake_numpy_namespace(self):
     # {{{ ArrayContext interface
 
     def clone(self):
-        return type(self)()
+        return type(self)(self.device)
 
     @overload
     def from_numpy(self, array: np.ndarray) -> Array:

From ce33f23983d3336dd280e75733be185e7d059f50 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 11 Feb 2025 13:17:43 -0800
Subject: [PATCH 25/34] remove loopy handling

---
 arraycontext/impl/cupy/__init__.py | 34 +++++++-----------------------
 1 file changed, 8 insertions(+), 26 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 17414d1b..50a81ee2 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -48,7 +48,6 @@
     ArrayOrContainerOrScalarT,
     ContainerOrScalarT,
     NumpyOrContainerOrScalar,
-    UntransformedCodeWarning,
 )
 
 
@@ -134,16 +133,10 @@ def call_loopy(
                 self,
                 t_unit: lp.TranslationUnit, **kwargs: Any
             ) -> dict[str, Array]:
-        t_unit = t_unit.copy(target=lp.ExecutableCTarget())
-        try:
-            executor = self._loopy_transform_cache[t_unit]
-        except KeyError:
-            executor = self.transform_loopy_program(t_unit).executor()
-            self._loopy_transform_cache[t_unit] = executor
-
-        _, result = executor(**kwargs)
-
-        return result
+        raise NotImplementedError(
+            "Calling loopy on CuPy arrays is not supported. Maybe rewrite"
+            " the loopy kernel as numpy-flavored array operations using"
+            " ArrayContext.np.")
 
     def freeze(self, array):
         import cupy as cp
@@ -164,21 +157,10 @@ def _thaw(ary):
     # }}}
 
     def transform_loopy_program(self, t_unit):
-        from warnings import warn
-        warn("Using the base "
-                f"{type(self).__name__}.transform_loopy_program "
-                "to transform a translation unit. "
-                "This is a no-op and will result in unoptimized C code for"
-                "the requested optimization, all in a single statement."
-                "This will work, but is unlikely to be performant."
-                f"Instead, subclass {type(self).__name__} and implement "
-                "the specific transform logic required to transform the program "
-                "for your package or application. Check higher-level packages "
-                "(e.g. meshmode), which may already have subclasses you may want "
-                "to build on.",
-                UntransformedCodeWarning, stacklevel=2)
-
-        return t_unit
+        raise NotImplementedError(
+            "Calling loopy on CuPy arrays is not supported. Maybe rewrite"
+            " the loopy kernel as numpy-flavored array operations using"
+            " ArrayContext.np.")
 
     def tag(self,
             tags: ToTagSetConvertible,

From cd2a36656583e5eca3c08e5cbe7e5ee814d4580b Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Wed, 12 Feb 2025 15:29:19 -0600
Subject: [PATCH 26/34] Revert "keep device in clone()"

This reverts commit 3ee28cc4d998a3379f3d601dafe744abeef2aa13.
---
 arraycontext/impl/cupy/__init__.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 50a81ee2..292c226e 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -74,8 +74,6 @@ def __init__(self, device: int | None = None) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
-        self.device = device
-
         if device is not None:
             import cupy as cp
             cp.cuda.runtime.setDevice(device)
@@ -89,7 +87,7 @@ def _get_fake_numpy_namespace(self):
     # {{{ ArrayContext interface
 
     def clone(self):
-        return type(self)(self.device)
+        return type(self)()
 
     @overload
     def from_numpy(self, array: np.ndarray) -> Array:

From 8e7e1f158a3f2f1ad8e1d2943512547e71e70a4e Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Wed, 12 Feb 2025 15:29:40 -0600
Subject: [PATCH 27/34] Revert "allow optional device selection"

This reverts commit ab8266d971bbb40467bb647d4fe504d67f72e513.
---
 arraycontext/impl/cupy/__init__.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 292c226e..d852cb90 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -70,14 +70,10 @@ class CupyArrayContext(ArrayContext):
 
     _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]
 
-    def __init__(self, device: int | None = None) -> None:
+    def __init__(self) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
-        if device is not None:
-            import cupy as cp
-            cp.cuda.runtime.setDevice(device)
-
     array_types = (CupyNonObjectArray,)
 
     def _get_fake_numpy_namespace(self):

From 326e1640438881991713ab20b4ede5da1a25dc7f Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 12:43:47 -0600
Subject: [PATCH 28/34] No need for separate method

Co-authored-by: Alex Fikl <alexfikl@gmail.com>
---
 arraycontext/impl/cupy/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index d852cb90..df977268 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -101,7 +101,7 @@ def from_numpy(self,
         def _from_numpy(ary):
             return cp.array(ary)
 
-        return with_array_context(rec_map_array_container(_from_numpy, array),
+        return with_array_context(rec_map_array_container(cp.array, array),
                                   actx=self)
 
     @overload
@@ -120,7 +120,7 @@ def to_numpy(self,
         def _to_numpy(ary):
             return cp.asnumpy(ary)
 
-        return with_array_context(rec_map_array_container(_to_numpy, array),
+        return with_array_context(rec_map_array_container(cp.asnumpy, array),
                                   actx=None)
 
     def call_loopy(

From 4d429f8bf47afe33ca6fd87221694af3b6992491 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 12:53:08 -0600
Subject: [PATCH 29/34] remove unneeded functions

---
 arraycontext/impl/cupy/__init__.py | 26 ++------------------------
 1 file changed, 2 insertions(+), 24 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index df977268..081ee38c 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -97,10 +97,6 @@ def from_numpy(self,
                    array: NumpyOrContainerOrScalar
                    ) -> ArrayOrContainerOrScalar:
         import cupy as cp
-
-        def _from_numpy(ary):
-            return cp.array(ary)
-
         return with_array_context(rec_map_array_container(cp.array, array),
                                   actx=self)
 
@@ -116,10 +112,6 @@ def to_numpy(self,
                  array: ArrayOrContainerOrScalar
                  ) -> NumpyOrContainerOrScalar:
         import cupy as cp
-
-        def _to_numpy(ary):
-            return cp.asnumpy(ary)
-
         return with_array_context(rec_map_array_container(cp.asnumpy, array),
                                   actx=None)
 
@@ -134,28 +126,14 @@ def call_loopy(
 
     def freeze(self, array):
         import cupy as cp
-
-        def _freeze(ary):
-            return cp.asnumpy(ary)
-
-        return with_array_context(rec_map_array_container(_freeze, array), actx=None)
+        return with_array_context(rec_map_array_container(cp.asnumpy, array), actx=None)
 
     def thaw(self, array):
         import cupy as cp
-
-        def _thaw(ary):
-            return cp.array(ary)
-
-        return with_array_context(rec_map_array_container(_thaw, array), actx=self)
+        return with_array_context(rec_map_array_container(cp.array, array), actx=self)
 
     # }}}
 
-    def transform_loopy_program(self, t_unit):
-        raise NotImplementedError(
-            "Calling loopy on CuPy arrays is not supported. Maybe rewrite"
-            " the loopy kernel as numpy-flavored array operations using"
-            " ArrayContext.np.")
-
     def tag(self,
             tags: ToTagSetConvertible,
             array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:

From 69e2133c61b8e441bd3ec38097af7db5daa95ef8 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 13:20:30 -0600
Subject: [PATCH 30/34] add a comment on non-blocking cp.asnumpy

---
 arraycontext/impl/cupy/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 081ee38c..cbc030ae 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -126,6 +126,8 @@ def call_loopy(
 
     def freeze(self, array):
         import cupy as cp
+        # Note that we could use a non-blocking version of cp.asnumpy here, but
+        # it appears to have very little impact on performance.
         return with_array_context(rec_map_array_container(cp.asnumpy, array), actx=None)
 
     def thaw(self, array):

From f83ee9c2a15b29fcfa16e64400ca18b7ea2a7390 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 13:43:52 -0600
Subject: [PATCH 31/34] add README links

---
 README.rst | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 1206093e..94a56a2d 100644
--- a/README.rst
+++ b/README.rst
@@ -15,8 +15,8 @@ GPU arrays? Deferred-evaluation arrays? Just plain ``numpy`` arrays? You'd like
 code to work with all of them? No problem! Comes with pre-made array context
 implementations for:
 
-- numpy
-- cupy
+- `numpy <https://numpy.org>`__
+- `cupy <https://cupy.dev>`__
 - `PyOpenCL <https://documen.tician.de/pyopencl/array.html>`__
 - `JAX <https://jax.readthedocs.io/en/latest/>`__
 - `Pytato <https://documen.tician.de/pytato>`__ (for lazy/deferred evaluation)

From fae85a43a8ac9fe9f946dc672b83183e27a3db40 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 13:44:32 -0600
Subject: [PATCH 32/34] fix typo

---
 README.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 94a56a2d..06800fc3 100644
--- a/README.rst
+++ b/README.rst
@@ -25,7 +25,7 @@ implementations for:
 - Profiling
 
 ``arraycontext`` started life as an array abstraction for use with the
-`meshmode <https://documen.tician.de/meshmode/>`__ unstrucuted discretization
+`meshmode <https://documen.tician.de/meshmode/>`__ unstructured discretization
 package.
 
 Distributed under the MIT license.

From d9ce8d5409cfd9c5dd7d0f8763cf7a4e41985560 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 17 Feb 2025 13:58:03 -0600
Subject: [PATCH 33/34] clean up loopy cache

Co-authored-by: Alex Fikl <alexfikl@gmail.com>
---
 arraycontext/impl/cupy/__init__.py | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index cbc030ae..4439e61f 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -62,17 +62,7 @@ class CupyNonObjectArray(metaclass=CupyNonObjectArrayMetaclass):
 
 
 class CupyArrayContext(ArrayContext):
-    """
-    An :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays.
-
-    .. automethod:: __init__
-    """
-
-    _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._loopy_transform_cache = {}
+    """An :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays."""
 
     array_types = (CupyNonObjectArray,)
 

From 22649a0115fc3603d9f84ecd09eb5e2fb66920b4 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Mon, 16 Jun 2025 17:26:09 -0500
Subject: [PATCH 34/34] bpr things

---
 arraycontext/container/__init__.py              | 8 +++++---
 arraycontext/container/traversal.py             | 2 +-
 arraycontext/context.py                         | 6 +++---
 arraycontext/impl/cupy/__init__.py              | 4 ++--
 arraycontext/impl/cupy/fake_numpy.py            | 2 +-
 arraycontext/impl/numpy/__init__.py             | 4 ++--
 arraycontext/impl/numpy/fake_numpy.py           | 2 +-
 arraycontext/impl/pyopencl/taggable_cl_array.py | 2 +-
 arraycontext/impl/pytato/compile.py             | 4 ++--
 9 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/arraycontext/container/__init__.py b/arraycontext/container/__init__.py
index 1ebc2cb2..a44f9cc9 100644
--- a/arraycontext/container/__init__.py
+++ b/arraycontext/container/__init__.py
@@ -94,6 +94,8 @@
 
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pymbolic.geometric_algebra import MultiVector
 
     from arraycontext import ArrayOrContainer
@@ -283,7 +285,7 @@ def get_container_context_opt(ary: ArrayContainer) -> ArrayContext | None:
 
 @serialize_container.register(np.ndarray)
 def _serialize_ndarray_container(
-        ary: numpy.ndarray) -> SerializedContainer:
+        ary: numpy.ndarray[Any, Any]) -> SerializedContainer:
     if ary.dtype.char != "O":
         raise NotAnArrayContainerError(
                 f"cannot serialize '{type(ary).__name__}' with dtype '{ary.dtype}'")
@@ -303,8 +305,8 @@ def _serialize_ndarray_container(
 @deserialize_container.register(np.ndarray)
 # https://github.com/python/mypy/issues/13040
 def _deserialize_ndarray_container(  # type: ignore[misc]
-        template: numpy.ndarray,
-        serialized: SerializedContainer) -> numpy.ndarray:
+        template: numpy.ndarray[Any, Any],
+        serialized: SerializedContainer) -> numpy.ndarray[Any, Any]:
     # disallow subclasses
     assert type(template) is np.ndarray
     assert template.dtype.char == "O"
diff --git a/arraycontext/container/traversal.py b/arraycontext/container/traversal.py
index 69d050f7..8620b7ec 100644
--- a/arraycontext/container/traversal.py
+++ b/arraycontext/container/traversal.py
@@ -917,7 +917,7 @@ def _flat_size(subary: ArrayOrContainer) -> Array | Integer:
 # {{{ numpy conversion
 
 def from_numpy(
-        ary: np.ndarray | ScalarLike,
+        ary: np.ndarray[Any, Any] | ScalarLike,
         actx: ArrayContext) -> ArrayOrContainerOrScalar:
     """Convert all :mod:`numpy` arrays in the :class:`~arraycontext.ArrayContainer`
     to the base array type of :class:`~arraycontext.ArrayContext`.
diff --git a/arraycontext/context.py b/arraycontext/context.py
index f751413c..61baf8ae 100644
--- a/arraycontext/context.py
+++ b/arraycontext/context.py
@@ -282,7 +282,7 @@ def __rtruediv__(self, other: Self | ScalarLike) -> Array: ...
 ContainerOrScalarT = TypeVar("ContainerOrScalarT", bound="ArrayContainer | ScalarLike")
 
 
-NumpyOrContainerOrScalar = Union[np.ndarray, "ArrayContainer", ScalarLike]
+NumpyOrContainerOrScalar = Union[np.ndarray[Any, Any], "ArrayContainer", ScalarLike]
 
 # }}}
 
@@ -358,7 +358,7 @@ def zeros(self,
         return self.np.zeros(shape, dtype)
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -379,7 +379,7 @@ def from_numpy(self,
         """
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
index 4439e61f..08e24bee 100644
--- a/arraycontext/impl/cupy/__init__.py
+++ b/arraycontext/impl/cupy/__init__.py
@@ -76,7 +76,7 @@ def clone(self):
         return type(self)()
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -91,7 +91,7 @@ def from_numpy(self,
                                   actx=self)
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload
diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py
index e45ca8a8..9ef85203 100644
--- a/arraycontext/impl/cupy/fake_numpy.py
+++ b/arraycontext/impl/cupy/fake_numpy.py
@@ -178,7 +178,7 @@ def linspace(self, *args, **kwargs):
         import cupy as cp
         return cp.linspace(*args, **kwargs)
 
-    def zeros_like(self, ary):
+    def zeros_like(self, ary):  # pyright: ignore[reportIncompatibleMethodOverride]
         import cupy as cp
         if isinstance(ary, int | float | complex):
             # Cupy does not support zeros_like with scalar arguments
diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py
index 6fcc88c7..d29fac7a 100644
--- a/arraycontext/impl/numpy/__init__.py
+++ b/arraycontext/impl/numpy/__init__.py
@@ -86,7 +86,7 @@ def clone(self):
         return type(self)()
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -99,7 +99,7 @@ def from_numpy(self,
         return array
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload
diff --git a/arraycontext/impl/numpy/fake_numpy.py b/arraycontext/impl/numpy/fake_numpy.py
index 582ccda9..b6702415 100644
--- a/arraycontext/impl/numpy/fake_numpy.py
+++ b/arraycontext/impl/numpy/fake_numpy.py
@@ -150,7 +150,7 @@ def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
                 return false_ary
             return np.logical_and.reduce(
                     [(true_ary if kx_i == ky_i else false_ary)
-                        and cast(np.ndarray, self.array_equal(x_i, y_i))
+                        and cast(np.ndarray, self.array_equal(x_i, y_i))  # pyright: ignore[reportMissingTypeArgument]
                         for (kx_i, x_i), (ky_i, y_i)
                         in zip(serialized_x, serialized_y, strict=True)],
                     initial=true_ary)
diff --git a/arraycontext/impl/pyopencl/taggable_cl_array.py b/arraycontext/impl/pyopencl/taggable_cl_array.py
index 79a108e2..60ea3f53 100644
--- a/arraycontext/impl/pyopencl/taggable_cl_array.py
+++ b/arraycontext/impl/pyopencl/taggable_cl_array.py
@@ -212,7 +212,7 @@ def zeros(
 
 def to_device(
             queue: cl.CommandQueue,
-            ary: np.ndarray[Any],
+            ary: np.ndarray[Any, Any],
             *, axes: tuple[Axis, ...] | None = None,
             tags: frozenset[Tag] = _EMPTY_TAG_SET,
             allocator: cla.Allocator | None = None,
diff --git a/arraycontext/impl/pytato/compile.py b/arraycontext/impl/pytato/compile.py
index e78c4e62..86e9bec1 100644
--- a/arraycontext/impl/pytato/compile.py
+++ b/arraycontext/impl/pytato/compile.py
@@ -100,12 +100,12 @@ def __hash__(self):
 
 @dataclass(frozen=True, eq=True)
 class ScalarInputDescriptor(AbstractInputDescriptor):
-    dtype: np.dtype
+    dtype: np.dtype[Any]
 
 
 @dataclass(frozen=True, eq=True)
 class LeafArrayDescriptor(AbstractInputDescriptor):
-    dtype: np.dtype
+    dtype: np.dtype[Any]
     shape: pt.array.ShapeType
 
 # }}}