doc/other.rst

 Other functionality
 ===================
 
+.. _metadata:
+
+Metadata ("tags") for Arrays and Array Axes
+-------------------------------------------
+
+.. automodule:: arraycontext.metadata
+
 :class:`~arraycontext.ArrayContext`-generating fixture for :mod:`pytest`
 ------------------------------------------------------------------------
 

doc/upload-docs.sh

 #! /bin/sh
 
-rsync --verbose --archive --delete _build/html/* doc-upload:doc/arraycontext
+rsync --verbose --archive --delete _build/html/ doc-upload:doc/arraycontext

pyproject.toml

+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "arraycontext"
+version = "2024.0"
+description = "Choose your favorite numpy-workalike"
+readme = "README.rst"
+license = { text = "MIT" }
+authors = [
+    { name = "Andreas Kloeckner", email = "inform@tiker.net" },
+]
+requires-python = ">=3.10"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Other Audience",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Natural Language :: English",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Utilities",
+]
+dependencies = [
+    "immutabledict>=4.1",
+    "numpy",
+    "pytools>=2024.1.3",
+    # for Self
+    "typing_extensions>=4",
+]
+
+[project.optional-dependencies]
+jax = [
+    "jax>=0.4",
+]
+pyopencl = [
+    "islpy>=2024.1",
+    "loopy>=2024.1",
+    "pyopencl>=2024.1",
+]
+pytato = [
+    "pytato>=2021.1",
+]
+test = [
+    "mypy",
+    "pytest",
+    "ruff",
+]
+
+[project.urls]
+Documentation = "https://documen.tician.de/arraycontext"
+Homepage = "https://github.com/inducer/arraycontext"
+
+[tool.ruff]
+preview = true
+
+[tool.ruff.lint]
+extend-select = [
+    "B",   # flake8-bugbear
+    "C",   # flake8-comprehensions
+    "E",   # pycodestyle
+    "F",   # pyflakes
+    "G",   # flake8-logging-format
+    "I",   # flake8-isort
+    "N",   # pep8-naming
+    "NPY", # numpy
+    "Q",   # flake8-quotes
+    "RUF", # ruff
+    "UP",  # pyupgrade
+    "W",   # pycodestyle
+    "SIM",
+]
+extend-ignore = [
+    "C90",   # McCabe complexity
+    "E221",  # multiple spaces before operator
+    "E226",  # missing whitespace around arithmetic operator
+    "E402",  # module-level import not at top of file
+]
+
+[tool.ruff.lint.flake8-quotes]
+docstring-quotes = "double"
+inline-quotes = "double"
+multiline-quotes = "double"
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+known-first-party = [
+    "jax",
+    "loopy",
+    "pymbolic",
+    "pyopencl",
+    "pytato",
+    "pytools",
+]
+known-local-folder = [
+    "arraycontext",
+    "testlib",
+]
+lines-after-imports = 2
+required-imports = ["from __future__ import annotations"]
+
+[tool.ruff.lint.per-file-ignores]
+"doc/conf.py" = ["I002"]
+# To avoid a requirement of array container definitions being someplace importable
+# from @dataclass_array_container.
+"test/test_utils.py" = ["I002"]
+
+[tool.mypy]
+python_version = "3.10"
+warn_unused_ignores = true
+# TODO: enable this
+# check_untyped_defs = true
+
+[[tool.mypy.overrides]]
+module = [
+    "islpy.*",
+    "loopy.*",
+    "meshmode.*",
+    "pymbolic",
+    "pymbolic.*",
+    "pyopencl.*",
+    "jax.*",
+]
+ignore_missing_imports = true
+
+[tool.typos.default]
+extend-ignore-re = [
+  "(?Rm)^.*(#|//)\\s*spellchecker:\\s*disable-line$"
+]

requirements.txt

 numpy
-recursivenodes
 git+https://github.com/inducer/pytools.git#egg=pytools
-git+https://github.com/inducer/gmsh_interop.git#egg=gmsh_interop
-git+https://github.com/inducer/pyvisfile.git#egg=pyvisfile
-git+https://github.com/inducer/modepy.git#egg=modepy
+git+https://github.com/inducer/pymbolic.git#egg=pymbolic
+
 git+https://github.com/inducer/pyopencl.git#egg=pyopencl
 git+https://github.com/inducer/islpy.git#egg=islpy
 
-# required by pytential, which is in turn needed for some tests
-git+https://github.com/inducer/pymbolic.git#egg=pymbolic
-
-# also depends on pymbolic, so should come after it
 git+https://github.com/inducer/loopy.git#egg=loopy
-
-# requires pymetis for tests for partition_mesh
 git+https://github.com/inducer/pytato.git#egg=pytato

run-mypy.sh

+#!/bin/bash
+
+python -m mypy --show-error-codes arraycontext examples test

run-pylint.sh

+#!/bin/bash
+
+set -o errexit -o nounset
+
+ci_support="https://gitlab.tiker.net/inducer/ci-support/raw/main"
+
+if [[ ! -f .pylintrc.yml ]]; then
+    curl -o .pylintrc.yml "${ci_support}/.pylintrc-default.yml"
+fi
+
+
+if [[ ! -f .run-pylint.py ]]; then
+    curl -L -o .run-pylint.py "${ci_support}/run-pylint.py"
+fi
+
+
+PYLINT_RUNNER_ARGS="--jobs=4 --yaml-rcfile=.pylintrc.yml"
+
+if [[ -f .pylintrc-local.yml ]]; then
+    PYLINT_RUNNER_ARGS+=" --yaml-rcfile=.pylintrc-local.yml"
+fi
+
+PYTHONWARNINGS=ignore python .run-pylint.py $PYLINT_RUNNER_ARGS $(basename $PWD) test/test_*.py examples "$@"

setup.cfg

-[flake8]
-ignore = E126,E127,E128,E123,E226,E241,E242,E265,W503,E402
-max-line-length=85
-
-inline-quotes = "
-docstring-quotes = """
-multiline-quotes = """
-
-# enable-flake8-bugbear

setup.py

-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-
-def main():
-    from setuptools import setup, find_packages
-
-    version_dict = {}
-    init_filename = "arraycontext/version.py"
-    exec(compile(open(init_filename, "r").read(), init_filename, "exec"),
-            version_dict)
-
-    setup(name="arraycontext",
-          version=version_dict["VERSION_TEXT"],
-          description="Choose your favorite numpy-workalike",
-          long_description=open("README.rst", "rt").read(),
-          author="Andreas Kloeckner",
-          author_email="inform@tiker.net",
-          license="MIT",
-          url="https://documen.tician.de/arraycontext",
-          classifiers=[
-              "Development Status :: 3 - Alpha",
-              "Intended Audience :: Developers",
-              "Intended Audience :: Other Audience",
-              "Intended Audience :: Science/Research",
-              "License :: OSI Approved :: MIT License",
-              "Natural Language :: English",
-              "Programming Language :: Python",
-              "Programming Language :: Python :: 3",
-              "Topic :: Scientific/Engineering",
-              "Topic :: Scientific/Engineering :: Information Analysis",
-              "Topic :: Scientific/Engineering :: Mathematics",
-              "Topic :: Scientific/Engineering :: Visualization",
-              "Topic :: Software Development :: Libraries",
-              "Topic :: Utilities",
-              ],
-
-          packages=find_packages(),
-          python_requires="~=3.6",
-          install_requires=[
-              "numpy",
-              "pytools>=2020.4.1",
-              "pytest>=2.3",
-              "loopy>=2019.1",
-              "dataclasses; python_version<='3.6'",
-              ],
-          )
-
-
-if __name__ == "__main__":
-    main()

test/test_arraycontext.py

+from __future__ import annotations
+
+
 __copyright__ = "Copyright (C) 2020-21 University of Illinois Board of Trustees"
 
 __license__ = """
@@ -20,182 +23,344 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
+import logging
 from dataclasses import dataclass
+from functools import partial
+
 import numpy as np
 import pytest
 
 from pytools.obj_array import make_obj_array
+from pytools.tag import Tag
 
 from arraycontext import (
-        ArrayContext,
-        dataclass_array_container, with_container_arithmetic,
-        serialize_container, deserialize_container,
-        freeze, thaw,
-        FirstAxisIsElementsTag)
-from arraycontext import (  # noqa: F401
-        pytest_generate_tests_for_pyopencl_array_context
-        as pytest_generate_tests)
+    BcastUntilActxArray,
+    EagerJAXArrayContext,
+    NumpyArrayContext,
+    PyOpenCLArrayContext,
+    PytatoPyOpenCLArrayContext,
+    dataclass_array_container,
+    pytest_generate_tests_for_array_contexts,
+    serialize_container,
+    tag_axes,
+    with_container_arithmetic,
+)
+from arraycontext.pytest import (
+    _PytestEagerJaxArrayContextFactory,
+    _PytestNumpyArrayContextFactory,
+    _PytestPyOpenCLArrayContextFactoryWithClass,
+    _PytestPytatoJaxArrayContextFactory,
+    _PytestPytatoPyOpenCLArrayContextFactory,
+)
+from testlib import DOFArray, MyContainer, MyContainerDOFBcast, Velocity2D
+
 
-import logging
 logger = logging.getLogger(__name__)
 
 
-# {{{ stand-in DOFArray implementation
+# {{{ array context fixture
+
+class _PyOpenCLArrayContextForTests(PyOpenCLArrayContext):
+    """Like :class:`PyOpenCLArrayContext`, but applies no program transformations
+    whatsoever. Only to be used for testing internal to :mod:`arraycontext`.
+    """
+
+    def transform_loopy_program(self, t_unit):
+        return t_unit
+
+
+class _PytatoPyOpenCLArrayContextForTests(PytatoPyOpenCLArrayContext):
+    """Like :class:`PytatoPyOpenCLArrayContext`, but applies no program
+    transformations whatsoever. Only to be used for testing internal to
+    :mod:`arraycontext`.
+    """
+
+    def transform_loopy_program(self, t_unit):
+        return t_unit
+
+
+class _PyOpenCLArrayContextWithHostScalarsForTestsFactory(
+        _PytestPyOpenCLArrayContextFactoryWithClass):
+    actx_class = _PyOpenCLArrayContextForTests
+
+
+class _PyOpenCLArrayContextForTestsFactory(
+        _PyOpenCLArrayContextWithHostScalarsForTestsFactory):
+    force_device_scalars = True
+
+
+class _PytatoPyOpenCLArrayContextForTestsFactory(
+        _PytestPytatoPyOpenCLArrayContextFactory):
+    actx_class = _PytatoPyOpenCLArrayContextForTests
+
+
+pytest_generate_tests = pytest_generate_tests_for_array_contexts([
+    _PyOpenCLArrayContextForTestsFactory,
+    _PyOpenCLArrayContextWithHostScalarsForTestsFactory,
+    _PytatoPyOpenCLArrayContextForTestsFactory,
+    _PytestEagerJaxArrayContextFactory,
+    _PytestPytatoJaxArrayContextFactory,
+    _PytestNumpyArrayContextFactory,
+    ])
+
+
+def _acf():
+    import pyopencl as cl
+
+    context = cl._csc()
+    queue = cl.CommandQueue(context)
+    return _PyOpenCLArrayContextForTests(queue, force_device_scalars=True)
+
+# }}}
+
+
+def _get_test_containers(actx, ambient_dim=2, shapes=50_000):
+    from numbers import Number
+
+    from testlib import DOFArray, MyContainer, MyContainerDOFBcast
+    if isinstance(shapes, Number | tuple):
+        shapes = [shapes]
+
+    x = DOFArray(actx, tuple(actx.from_numpy(randn(shape, np.float64))
+        for shape in shapes))
+
+    # pylint: disable=unexpected-keyword-arg, no-value-for-parameter
+    dataclass_of_dofs = MyContainer(
+            name="container",
+            mass=x,
+            momentum=make_obj_array([x] * ambient_dim),
+            enthalpy=x)
+
+    # pylint: disable=unexpected-keyword-arg, no-value-for-parameter
+    bcast_dataclass_of_dofs = MyContainerDOFBcast(
+            name="container",
+            mass=x,
+            momentum=make_obj_array([x] * ambient_dim),
+            enthalpy=x)
+
+    ary_dof = x
+    ary_of_dofs = make_obj_array([x] * ambient_dim)
+    mat_of_dofs = np.empty((ambient_dim, ambient_dim), dtype=object)
+    for i in np.ndindex(mat_of_dofs.shape):
+        mat_of_dofs[i] = x
+
+    return (ary_dof, ary_of_dofs, mat_of_dofs, dataclass_of_dofs,
+            bcast_dataclass_of_dofs)
 
-@with_container_arithmetic(
-        bcast_obj_array=True,
-        bcast_numpy_array=True,
-        rel_comparison=True)
-class DOFArray:
-    def __init__(self, actx, data):
-        if not (actx is None or isinstance(actx, ArrayContext)):
-            raise TypeError("actx must be of type ArrayContext")
 
-        if not isinstance(data, tuple):
-            raise TypeError("'data' argument must be a tuple")
+# {{{ assert_close_to_numpy*
 
-        self.array_context = actx
-        self.data = data
+def randn(shape, dtype):
+    rng = np.random.default_rng()
+    dtype = np.dtype(dtype)
+
+    ashape = 1 if shape == 0 else shape
+
+    if dtype.kind == "c":
+        dtype = np.dtype(f"<f{dtype.itemsize // 2}")
+        r = rng.standard_normal(ashape, dtype) \
+            + 1j * rng.standard_normal(ashape, dtype)
+    elif dtype.kind == "f":
+        r = rng.standard_normal(ashape, dtype)
+    elif dtype.kind == "i":
+        r = rng.integers(0, 512, ashape, dtype)
+    else:
+        raise TypeError(dtype.kind)
 
-    __array_priority__ = 10
+    if shape == 0:
+        return np.array(r[0])
 
-    def __len__(self):
-        return len(self.data)
+    return r
 
-    def __getitem__(self, i):
-        return self.data[i]
 
-    @classmethod
-    def _serialize_init_arrays_code(cls, instance_name):
-        return {"_":
-                (f"{instance_name}_i", f"{instance_name}")}
+def assert_close_to_numpy(actx, op, args):
+    assert np.allclose(
+            actx.to_numpy(
+                op(actx.np, *[
+                    actx.from_numpy(arg) if isinstance(arg, np.ndarray) else arg
+                    for arg in args])),
+            op(np, *args))
 
-    @classmethod
-    def _deserialize_init_arrays_code(cls, template_instance_name, args):
-        (_, arg), = args.items()
-        # Why tuple([...])? https://stackoverflow.com/a/48592299
-        return (f"{template_instance_name}.array_context, tuple([{arg}])")
 
-    @property
-    def real(self):
-        return DOFArray(self.array_context, tuple([subary.real for subary in self]))
+def assert_close_to_numpy_in_containers(actx, op, args):
+    assert_close_to_numpy(actx, op, args)
 
-    @property
-    def imag(self):
-        return DOFArray(self.array_context, tuple([subary.imag for subary in self]))
+    ref_result = op(np, *args)
 
+    # {{{ test DOFArrays
 
-@serialize_container.register(DOFArray)
-def _serialize_dof_container(ary: DOFArray):
-    return enumerate(ary.data)
+    dofarray_args = [
+            DOFArray(actx, (actx.from_numpy(arg),))
+            if isinstance(arg, np.ndarray) else arg
+            for arg in args]
 
+    actx_result = op(actx.np, *dofarray_args)
+    if isinstance(actx_result, DOFArray):
+        actx_result = actx_result[0]
 
-@deserialize_container.register(DOFArray)
-def _deserialize_dof_container(
-        template, iterable):
-    def _raise_index_inconsistency(i, stream_i):
-        raise ValueError(
-                "out-of-sequence indices supplied in DOFArray deserialization "
-                f"(expected {i}, received {stream_i})")
+    assert np.allclose(actx.to_numpy(actx_result), ref_result)
 
-    return type(template)(
-            template.array_context,
-            data=tuple(
-                v if i == stream_i else _raise_index_inconsistency(i, stream_i)
-                for i, (stream_i, v) in enumerate(iterable)))
+    # }}}
 
+    # {{{ test object arrays of DOFArrays
 
-@freeze.register(DOFArray)
-def _freeze_dofarray(ary, actx=None):
-    assert actx is None
-    return type(ary)(
-        None,
-        tuple(ary.array_context.freeze(subary) for subary in ary.data))
+    obj_array_args = [
+            make_obj_array([arg]) if isinstance(arg, DOFArray) else arg
+            for arg in dofarray_args]
 
+    obj_array_result = op(actx.np, *obj_array_args)
+    if isinstance(obj_array_result, np.ndarray):
+        obj_array_result = obj_array_result[0][0]
 
-@thaw.register(DOFArray)
-def _thaw_dofarray(ary, actx):
-    if ary.array_context is not None:
-        raise ValueError("cannot thaw DOFArray that already has an array context")
+    assert np.allclose(actx.to_numpy(obj_array_result), ref_result)
 
-    return type(ary)(
-        actx,
-        tuple(actx.thaw(subary) for subary in ary.data))
+    # }}}
 
 # }}}
 
 
-def test_array_context_np_workalike(actx_factory):
+# {{{ np.function same as numpy
+
+@pytest.mark.parametrize(("sym_name", "n_args", "dtype"), [
+            # float only
+            ("arctan2", 2, np.float64),
+            ("minimum", 2, np.float64),
+            ("maximum", 2, np.float64),
+            ("where", 3, np.float64),
+            ("min", 1, np.float64),
+            ("max", 1, np.float64),
+            ("any", 1, np.float64),
+            ("all", 1, np.float64),
+            ("arctan", 1, np.float64),
+
+            # float + complex
+            ("sin", 1, np.float64),
+            ("sin", 1, np.complex128),
+            ("exp", 1, np.float64),
+            ("exp", 1, np.complex128),
+            ("conj", 1, np.float64),
+            ("conj", 1, np.complex128),
+            ("vdot", 2, np.float64),
+            ("vdot", 2, np.complex128),
+            ("abs", 1, np.float64),
+            ("abs", 1, np.complex128),
+            ("sum", 1, np.float64),
+            ("sum", 1, np.complex64),
+            ("isnan", 1, np.float64),
+            ])
+def test_array_context_np_workalike(actx_factory, sym_name, n_args, dtype):
     actx = actx_factory()
+    if not hasattr(actx.np, sym_name):
+        pytest.skip(f"'{sym_name}' not implemented on '{type(actx).__name__}'")
 
-    ndofs = 5000
+    ndofs = 512
+    args = [randn(ndofs, dtype) for i in range(n_args)]
 
-    for sym_name, n_args in [
-            ("sin", 1),
-            ("exp", 1),
-            ("arctan2", 2),
-            ("minimum", 2),
-            ("maximum", 2),
-            ("where", 3),
-            ("conj", 1),
-            # ("empty_like", 1),    # NOTE: fails np.allclose, obviously
-            ("zeros_like", 1),
-            ("ones_like", 1),
-            ]:
-        args = [np.random.randn(ndofs) for i in range(n_args)]
-        ref_result = getattr(np, sym_name)(*args)
+    c_to_numpy_arc_functions = {
+            "atan": "arctan",
+            "atan2": "arctan2",
+            }
+
+    def evaluate(np_, *args_):
+        func = getattr(np_, sym_name,
+                getattr(np_, c_to_numpy_arc_functions.get(sym_name, sym_name)))
 
-        # {{{ test cl.Arrays
+        return func(*args_)
 
-        actx_args = [actx.from_numpy(arg) for arg in args]
-        actx_result = actx.to_numpy(getattr(actx.np, sym_name)(*actx_args))
+    assert_close_to_numpy_in_containers(actx, evaluate, args)
 
-        assert np.allclose(actx_result, ref_result)
+    if sym_name not in ["where", "min", "max", "any", "all", "conj", "vdot", "sum"]:
+        # Scalar arguments are supported.
+        args = [randn(0, dtype)[()] for i in range(n_args)]
+        assert_close_to_numpy(actx, evaluate, args)
 
-        # }}}
 
-        # {{{ test DOFArrays
+@pytest.mark.parametrize(("sym_name", "n_args", "dtype"), [
+            ("zeros_like", 1, np.float64),
+            ("zeros_like", 1, np.complex128),
+            ("ones_like", 1, np.float64),
+            ("ones_like", 1, np.complex128),
+            ])
+def test_array_context_np_like(actx_factory, sym_name, n_args, dtype):
+    actx = actx_factory()
 
-        actx_args = [DOFArray(actx, (arg,)) for arg in actx_args]
-        actx_result = actx.to_numpy(
-                getattr(actx.np, sym_name)(*actx_args)[0])
+    ndofs = 512
+    args = [randn(ndofs, dtype) for i in range(n_args)]
+    assert_close_to_numpy(
+            actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
+
+    for c in (42.0, *_get_test_containers(actx)):
+        result = getattr(actx.np, sym_name)(c)
+        result = actx.thaw(actx.freeze(result))
+
+        if sym_name == "zeros_like":
+            if np.isscalar(result):
+                assert result == 0.0
+            else:
+                assert actx.to_numpy(actx.np.all(actx.np.equal(result, 0.0)))
+        elif sym_name == "ones_like":
+            if np.isscalar(result):
+                assert result == 1.0
+            else:
+                assert actx.to_numpy(actx.np.all(actx.np.equal(result, 1.0)))
+        else:
+            raise ValueError(f"unknown method: '{sym_name}'")
 
-        assert np.allclose(actx_result, ref_result)
+# }}}
 
-        # }}}
 
-        # {{{ test object arrays
+# {{{ array manipulations
 
-        if sym_name in ("empty_like", "zeros_like", "ones_like"):
-            continue
+def test_actx_stack(actx_factory):
+    rng = np.random.default_rng()
 
-        obj_array_args = [make_obj_array([arg]) for arg in actx_args]
-        obj_array_result = actx.to_numpy(
-                getattr(actx.np, sym_name)(*obj_array_args)[0][0])
+    actx = actx_factory()
 
-        assert np.allclose(obj_array_result, ref_result)
+    ndofs = 5000
+    args = [rng.normal(size=ndofs) for i in range(10)]
 
-        # }}}
+    assert_close_to_numpy_in_containers(
+            actx, lambda _np, *_args: _np.stack(_args), args)
 
 
 def test_actx_concatenate(actx_factory):
+    rng = np.random.default_rng()
     actx = actx_factory()
 
     ndofs = 5000
-    args = [np.random.randn(ndofs) for i in range(10)]
-    ref_result = np.concatenate(args)
+    args = [rng.normal(size=ndofs) for i in range(10)]
 
-    # {{{ test cl.Arrays
+    assert_close_to_numpy(
+            actx, lambda _np, *_args: _np.concatenate(_args), args)
 
-    actx_args = [actx.from_numpy(arg) for arg in args]
-    actx_result = actx.to_numpy(actx.np.concatenate(actx_args))
 
-    assert np.allclose(actx_result, ref_result)
+def test_actx_reshape(actx_factory):
+    rng = np.random.default_rng()
+    actx = actx_factory()
+
+    for new_shape in [(3, 2), (3, -1), (6,), (-1,)]:
+        assert_close_to_numpy(
+                actx, lambda _np, *_args: _np.reshape(*_args),
+                (rng.normal(size=(2, 3)), new_shape))
 
-    # }}}
 
+def test_actx_ravel(actx_factory):
+    from numpy.random import default_rng
+    actx = actx_factory()
+    rng = default_rng()
+    ndim = rng.integers(low=1, high=6)
+    shape = tuple(rng.integers(2, 7, ndim))
+
+    assert_close_to_numpy(actx, lambda _np, ary: _np.ravel(ary),
+                          (rng.random(shape),))
+
+# }}}
+
+
+# {{{ arithmetic same as numpy
 
 def test_dof_array_arithmetic_same_as_numpy(actx_factory):
+    rng = np.random.default_rng()
     actx = actx_factory()
 
     ndofs = 50_000
@@ -204,11 +369,12 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
         return ary.real
 
     def get_imag(ary):
-        return ary.real
+        return ary.imag
 
     import operator
+    from random import randrange, uniform
+
     from pytools import generate_nonnegative_integer_tuples_below as gnitb
-    from random import uniform, randrange
     for op_func, n_args, use_integers in [
             (operator.add, 2, False),
             (operator.sub, 2, False),
@@ -216,12 +382,12 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
             (operator.truediv, 2, False),
             (operator.pow, 2, False),
             # FIXME pyopencl.Array doesn't do mod.
-            #(operator.mod, 2, True),
-            #(operator.mod, 2, False),
-            #(operator.imod, 2, True),
-            #(operator.imod, 2, False),
+            # (operator.mod, 2, True),
+            # (operator.mod, 2, False),
+            # (operator.imod, 2, True),
+            # (operator.imod, 2, False),
             # FIXME: Two outputs
-            #(divmod, 2, False),
+            # (divmod, 2, False),
 
             (operator.iadd, 2, False),
             (operator.isub, 2, False),
@@ -262,10 +428,23 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
                 # can't do in place operations with a scalar lhs
                 continue
 
+            if op_func == operator.ge:
+                op_func_actx = actx.np.greater_equal
+            elif op_func == operator.lt:
+                op_func_actx = actx.np.less
+            elif op_func == operator.gt:
+                op_func_actx = actx.np.greater
+            elif op_func == operator.eq:
+                op_func_actx = actx.np.equal
+            elif op_func == operator.ne:
+                op_func_actx = actx.np.not_equal
+            else:
+                op_func_actx = op_func
+
             args = [
-                    (0.5+np.random.rand(ndofs)
+                    (0.5+rng.uniform(size=ndofs)
                         if not use_integers else
-                        np.random.randint(3, 200, ndofs))
+                        rng.integers(3, 200, size=ndofs))
 
                     if is_array_flag else
                     (uniform(0.5, 2)
@@ -290,7 +469,7 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
                     if isinstance(arg, np.ndarray) else arg
                     for arg in args]
 
-            actx_result = actx.to_numpy(op_func(*actx_args)[0])
+            actx_result = actx.to_numpy(op_func_actx(*actx_args)[0])
 
             assert np.allclose(actx_result, ref_result)
 
@@ -322,24 +501,86 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
                         for arg in actx_args]
 
                 obj_array_result = actx.to_numpy(
-                        op_func(*obj_array_args)[0][0])
+                        op_func_actx(*obj_array_args)[0][0])
 
                 assert np.allclose(obj_array_result, ref_result)
 
             # }}}
 
+# }}}
+
+
+# {{{ reductions same as numpy
 
-def test_dof_array_reductions_same_as_numpy(actx_factory):
+@pytest.mark.parametrize("op", ["sum", "min", "max"])
+def test_reductions_same_as_numpy(actx_factory, op):
+    rng = np.random.default_rng()
     actx = actx_factory()
 
-    from numbers import Number
-    for name in ["sum", "min", "max"]:
-        ary = np.random.randn(3000)
-        np_red = getattr(np, name)(ary)
-        actx_red = getattr(actx.np, name)(actx.from_numpy(ary))
+    ary = rng.normal(size=3000)
+    np_red = getattr(np, op)(ary)
+    actx_red = getattr(actx.np, op)(actx.from_numpy(ary))
+    actx_red = actx.to_numpy(actx_red)
+
+    assert actx_red.shape == ()
+
+    assert np.allclose(np_red, actx_red)
+
+
+@pytest.mark.parametrize("sym_name", ["any", "all"])
+def test_any_all_same_as_numpy(actx_factory, sym_name):
+    actx = actx_factory()
+    if not hasattr(actx.np, sym_name):
+        pytest.skip(f"'{sym_name}' not implemented on '{type(actx).__name__}'")
+
+    rng = np.random.default_rng()
+    ary_any = rng.integers(0, 2, 512)
+    ary_all = np.ones(512)
+
+    assert_close_to_numpy_in_containers(actx,
+                lambda _np, *_args: getattr(_np, sym_name)(*_args), [ary_any])
+    assert_close_to_numpy_in_containers(actx,
+                lambda _np, *_args: getattr(_np, sym_name)(*_args), [ary_all])
+    assert_close_to_numpy_in_containers(actx,
+                lambda _np, *_args: getattr(_np, sym_name)(*_args), [1 - ary_all])
+
+
+def test_array_equal(actx_factory):
+    actx = actx_factory()
+
+    sym_name = "array_equal"
+    if not hasattr(actx.np, sym_name):
+        pytest.skip(f"'{sym_name}' not implemented on '{type(actx).__name__}'")
 
-        assert isinstance(actx_red, Number)
-        assert np.allclose(np_red, actx_red)
+    rng = np.random.default_rng()
+    ary = rng.integers(0, 2, 512)
+    ary_copy = ary.copy()
+    ary_diff_values = np.ones(512)
+    ary_diff_shape = np.ones(511)
+    ary_diff_type = DOFArray(actx, (np.ones(512),))
+
+    # Equal
+    assert_close_to_numpy_in_containers(actx,
+        lambda _np, *_args: getattr(_np, sym_name)(*_args), [ary, ary_copy])
+
+    # Different values
+    assert_close_to_numpy_in_containers(actx,
+        lambda _np, *_args: getattr(_np, sym_name)(*_args), [ary, ary_diff_values])
+
+    # Different shapes
+    assert_close_to_numpy_in_containers(actx,
+        lambda _np, *_args: getattr(_np, sym_name)(*_args), [ary, ary_diff_shape])
+
+    # Different types
+    assert not actx.to_numpy(actx.np.array_equal(ary, ary_diff_type))
+
+    # Empty
+    ary_empty = np.empty((5, 0), dtype=object)
+    ary_empty_copy = ary_empty.copy()
+    assert actx.to_numpy(actx.np.array_equal(ary_empty, ary_empty_copy))
+
+
+# }}}
 
 
 # {{{ test array context einsum
@@ -351,10 +592,10 @@ def test_dof_array_reductions_same_as_numpy(actx_factory):
 ])
 def test_array_context_einsum_array_manipulation(actx_factory, spec):
     actx = actx_factory()
+    rng = np.random.default_rng()
 
-    mat = actx.from_numpy(np.random.randn(10, 10))
-    res = actx.to_numpy(actx.einsum(spec, mat,
-                                    tagged=(FirstAxisIsElementsTag())))
+    mat = actx.from_numpy(rng.normal(size=(10, 10)))
+    res = actx.to_numpy(actx.einsum(spec, mat))
     ans = np.einsum(spec, actx.to_numpy(mat))
     assert np.allclose(res, ans)
 
@@ -366,11 +607,11 @@ def test_array_context_einsum_array_manipulation(actx_factory, spec):
 ])
 def test_array_context_einsum_array_matmatprods(actx_factory, spec):
     actx = actx_factory()
+    rng = np.random.default_rng()
 
-    mat_a = actx.from_numpy(np.random.randn(5, 5))
-    mat_b = actx.from_numpy(np.random.randn(5, 5))
-    res = actx.to_numpy(actx.einsum(spec, mat_a, mat_b,
-                                    tagged=(FirstAxisIsElementsTag())))
+    mat_a = actx.from_numpy(rng.normal(size=(5, 5)))
+    mat_b = actx.from_numpy(rng.normal(size=(5, 5)))
+    res = actx.to_numpy(actx.einsum(spec, mat_a, mat_b))
     ans = np.einsum(spec, actx.to_numpy(mat_a), actx.to_numpy(mat_b))
     assert np.allclose(res, ans)
 
@@ -380,12 +621,12 @@ def test_array_context_einsum_array_matmatprods(actx_factory, spec):
 ])
 def test_array_context_einsum_array_tripleprod(actx_factory, spec):
     actx = actx_factory()
+    rng = np.random.default_rng()
 
-    mat_a = actx.from_numpy(np.random.randn(7, 5))
-    mat_b = actx.from_numpy(np.random.randn(5, 7))
-    vec = actx.from_numpy(np.random.randn(7))
-    res = actx.to_numpy(actx.einsum(spec, mat_a, mat_b, vec,
-                                    tagged=(FirstAxisIsElementsTag())))
+    mat_a = actx.from_numpy(rng.normal(size=(7, 5)))
+    mat_b = actx.from_numpy(rng.normal(size=(5, 7)))
+    vec = actx.from_numpy(rng.normal(size=(7)))
+    res = actx.to_numpy(actx.einsum(spec, mat_a, mat_b, vec))
     ans = np.einsum(spec,
                     actx.to_numpy(mat_a),
                     actx.to_numpy(mat_b),
@@ -395,49 +636,111 @@ def test_array_context_einsum_array_tripleprod(actx_factory, spec):
 # }}}
 
 
-# {{{ test array container
+# {{{ array container classes for test
 
-@with_container_arithmetic(bcast_obj_array=False, rel_comparison=True)
-@dataclass_array_container
-@dataclass(frozen=True)
-class MyContainer:
-    name: str
-    mass: DOFArray
-    momentum: np.ndarray
-    enthalpy: DOFArray
 
-    @property
-    def array_context(self):
-        return self.mass.array_context
+def test_container_map_on_device_scalar(actx_factory):
+    actx = actx_factory()
 
+    expected_sizes = [1, 2, 4, 4, 4]
+    arys = _get_test_containers(actx, shapes=0)
+    arys += (np.pi,)
 
-def _get_test_containers(actx, ambient_dim=2):
-    x = DOFArray(actx, (actx.from_numpy(np.random.randn(50_000)),))
+    from arraycontext import (
+        map_array_container,
+        map_reduce_array_container,
+        rec_map_array_container,
+        rec_map_reduce_array_container,
+    )
 
-    # pylint: disable=unexpected-keyword-arg, no-value-for-parameter
-    dataclass_of_dofs = MyContainer(
-            name="container",
-            mass=x,
-            momentum=make_obj_array([x, x]),
-            enthalpy=x)
+    for size, ary in zip(expected_sizes, arys[:-1], strict=True):
+        result = map_array_container(lambda x: x, ary)
+        assert actx.to_numpy(actx.np.array_equal(result, ary))
+        result = rec_map_array_container(lambda x: x, ary)
+        assert actx.to_numpy(actx.np.array_equal(result, ary))
 
-    ary_dof = x
-    ary_of_dofs = make_obj_array([x, x, x])
-    mat_of_dofs = np.empty((2, 2), dtype=object)
-    for i in np.ndindex(mat_of_dofs.shape):
-        mat_of_dofs[i] = x
+        result = map_reduce_array_container(sum, np.size, ary)
+        assert result == size
+        result = rec_map_reduce_array_container(sum, np.size, ary)
+        assert result == size
+
+
+def test_container_map(actx_factory):
+    actx = actx_factory()
+    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs, _bcast_dc_of_dofs = \
+            _get_test_containers(actx)
+
+    # {{{ check
 
-    return ary_dof, ary_of_dofs, mat_of_dofs, dataclass_of_dofs
+    def _check_allclose(f, arg1, arg2, atol=2.0e-14):
+        from arraycontext import NotAnArrayContainerError
+        try:
+            arg1_iterable = serialize_container(arg1)
+            arg2_iterable = serialize_container(arg2)
+        except NotAnArrayContainerError:
+            assert np.linalg.norm(actx.to_numpy(f(arg1) - arg2)) < atol
+        else:
+            arg1_subarrays = [
+                subarray for _, subarray in arg1_iterable]
+            arg2_subarrays = [
+                subarray for _, subarray in arg2_iterable]
+            for subarray1, subarray2 in zip(arg1_subarrays, arg2_subarrays,
+                                            strict=True):
+                _check_allclose(f, subarray1, subarray2)
+
+    def func(x):
+        return x + 1
+
+    from arraycontext import rec_map_array_container
+    result = rec_map_array_container(func, 1)
+    assert result == 2
+
+    for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
+        result = rec_map_array_container(func, ary)
+        _check_allclose(func, ary, result)
+
+    from arraycontext import mapped_over_array_containers
+
+    @mapped_over_array_containers
+    def mapped_func(x):
+        return func(x)
+
+    for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
+        result = mapped_func(ary)
+        _check_allclose(func, ary, result)
+
+    @mapped_over_array_containers(leaf_class=DOFArray)
+    def check_leaf(x):
+        assert isinstance(x, DOFArray)
+
+    for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
+        check_leaf(ary)
+
+    # }}}
 
 
 def test_container_multimap(actx_factory):
     actx = actx_factory()
-    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs = _get_test_containers(actx)
+    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs, _bcast_dc_of_dofs = \
+            _get_test_containers(actx)
 
     # {{{ check
 
-    def _check_allclose(f, arg1, arg2, atol=1.0e-14):
-        assert np.linalg.norm((f(arg1) - arg2).get()) < atol
+    def _check_allclose(f, arg1, arg2, atol=2.0e-14):
+        from arraycontext import NotAnArrayContainerError
+        try:
+            arg1_iterable = serialize_container(arg1)
+            arg2_iterable = serialize_container(arg2)
+        except NotAnArrayContainerError:
+            assert np.linalg.norm(actx.to_numpy(f(arg1) - arg2)) < atol
+        else:
+            arg1_subarrays = [
+                subarray for _, subarray in arg1_iterable]
+            arg2_subarrays = [
+                subarray for _, subarray in arg2_iterable]
+            for subarray1, subarray2 in zip(arg1_subarrays, arg2_subarrays,
+                                            strict=True):
+                _check_allclose(f, subarray1, subarray2)
 
     def func_all_scalar(x, y):
         return x + y
@@ -452,17 +755,30 @@ def test_container_multimap(actx_factory):
     result = rec_multimap_array_container(func_all_scalar, 1, 2)
     assert result == 3
 
-    from functools import partial
     for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
         result = rec_multimap_array_container(func_first_scalar, 1, ary)
-        rec_multimap_array_container(
-                partial(_check_allclose, lambda x: 1 + x),
-                ary, result)
+        _check_allclose(lambda x: 1 + x, ary, result)
 
         result = rec_multimap_array_container(func_multiple_scalar, 2, ary, 2, ary)
-        rec_multimap_array_container(
-                partial(_check_allclose, lambda x: 4 * x),
-                ary, result)
+        _check_allclose(lambda x: 4 * x, ary, result)
+
+    from arraycontext import multimapped_over_array_containers
+
+    @multimapped_over_array_containers
+    def mapped_func(a, subary1, b, subary2):
+        return func_multiple_scalar(a, subary1, b, subary2)
+
+    for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
+        result = mapped_func(2, ary, 2, ary)
+        _check_allclose(lambda x: 4 * x, ary, result)
+
+    @multimapped_over_array_containers(leaf_class=DOFArray)
+    def check_leaf(a, subary1, b, subary2):
+        assert isinstance(subary1, DOFArray)
+        assert isinstance(subary2, DOFArray)
+
+    for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
+        check_leaf(2, ary, 2, ary)
 
     with pytest.raises(AssertionError):
         rec_multimap_array_container(func_multiple_scalar, 2, ary_dof, 2, dc_of_dofs)
@@ -472,14 +788,14 @@ def test_container_multimap(actx_factory):
 
 def test_container_arithmetic(actx_factory):
     actx = actx_factory()
-    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs = _get_test_containers(actx)
+    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs, bcast_dc_of_dofs = \
+            _get_test_containers(actx)
 
     # {{{ check
 
-    def _check_allclose(f, arg1, arg2, atol=1.0e-14):
-        assert np.linalg.norm((f(arg1) - arg2).get()) < atol
+    def _check_allclose(f, arg1, arg2, atol=5.0e-14):
+        assert np.linalg.norm(actx.to_numpy(f(arg1) - arg2)) < atol
 
-    from functools import partial
     from arraycontext import rec_multimap_array_container
     for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
         rec_multimap_array_container(
@@ -492,33 +808,78 @@ def test_container_arithmetic(actx_factory):
     with pytest.raises(TypeError):
         ary_of_dofs + dc_of_dofs
 
+    with pytest.raises(TypeError):
+        dc_of_dofs + ary_of_dofs
+
+    with pytest.raises(TypeError):
+        ary_dof + dc_of_dofs
+
+    with pytest.raises(TypeError):
+        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
+
+    mock_gradient = MyContainerDOFBcast(
+            name="yo",
+            mass=ary_of_dofs,
+            momentum=mat_of_dofs,
+            enthalpy=ary_of_dofs)
+
+    grad_matvec_result = mock_gradient @ ary_of_dofs
+    assert isinstance(grad_matvec_result.mass, DOFArray)
+    assert grad_matvec_result.momentum.shape == ary_of_dofs.shape
+
+    assert actx.to_numpy(actx.np.linalg.norm(
+        grad_matvec_result.mass - sum(ary_of_dofs**2)
+        )) < 1e-8
+
     # }}}
 
 
 def test_container_freeze_thaw(actx_factory):
     actx = actx_factory()
-    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs = _get_test_containers(actx)
+    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs, _bcast_dc_of_dofs = \
+            _get_test_containers(actx)
 
     # {{{ check
 
-    from arraycontext import get_container_context
-    from arraycontext import get_container_context_recursively
+    from arraycontext import (
+        get_container_context_opt,
+        get_container_context_recursively_opt,
+    )
 
-    assert get_container_context(ary_of_dofs) is None
-    assert get_container_context(mat_of_dofs) is None
-    assert get_container_context(ary_dof) is actx
-    assert get_container_context(dc_of_dofs) is actx
+    assert get_container_context_opt(ary_of_dofs) is None
+    assert get_container_context_opt(mat_of_dofs) is None
+    assert get_container_context_opt(ary_dof) is actx
+    assert get_container_context_opt(dc_of_dofs) is actx
 
-    assert get_container_context_recursively(ary_of_dofs) is actx
-    assert get_container_context_recursively(mat_of_dofs) is actx
+    assert get_container_context_recursively_opt(ary_of_dofs) is actx
+    assert get_container_context_recursively_opt(mat_of_dofs) is actx
 
     for ary in [ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs]:
-        frozen_ary = freeze(ary)
-        thawed_ary = thaw(frozen_ary, actx)
-        frozen_ary = freeze(thawed_ary)
+        frozen_ary = actx.freeze(ary)
+        thawed_ary = actx.thaw(frozen_ary)
+        frozen_ary = actx.freeze(thawed_ary)
+
+        assert get_container_context_recursively_opt(frozen_ary) is None
+        assert get_container_context_recursively_opt(thawed_ary) is actx
+
+    actx2 = actx.clone()
+
+    ary_dof_frozen = actx.freeze(ary_dof)
+    with pytest.raises(ValueError) as exc_info:
+        ary_dof + ary_dof_frozen
 
-        assert get_container_context_recursively(frozen_ary) is None
-        assert get_container_context_recursively(thawed_ary) is actx
+    assert "frozen" in str(exc_info.value)
+
+    ary_dof_2 = actx2.thaw(actx.freeze(ary_dof))
+
+    with pytest.raises(ValueError):
+        ary_dof + ary_dof_2
 
     # }}}
 
@@ -527,8 +888,6 @@ def test_container_freeze_thaw(actx_factory):
 def test_container_norm(actx_factory, ord):
     actx = actx_factory()
 
-    ary_dof, ary_of_dofs, mat_of_dofs, dc_of_dofs = _get_test_containers(actx)
-
     from pytools.obj_array import make_obj_array
     c = MyContainer(name="hey", mass=1, momentum=make_obj_array([2, 3]), enthalpy=5)
     n1 = actx.np.linalg.norm(make_obj_array([c, c]), ord)
@@ -539,6 +898,631 @@ def test_container_norm(actx_factory, ord):
 # }}}
 
 
+# {{{ test flatten and unflatten
+
+@pytest.mark.parametrize("shapes", [
+    0,                          # tests device scalars when flattening
+    512,
+    [(128, 67)],
+    [(127, 67), (18, 0)],       # tests 0-sized arrays
+    [(64, 7), (154, 12)]
+    ])
+def test_flatten_array_container(actx_factory, shapes):
+    actx = actx_factory()
+
+    from arraycontext import flatten, unflatten
+    arys = _get_test_containers(actx, shapes=shapes)
+
+    for ary in arys:
+        flat = flatten(ary, actx)
+        assert flat.ndim == 1
+
+        ary_roundtrip = unflatten(ary, flat, actx)
+
+        from arraycontext import rec_multimap_reduce_array_container
+        assert rec_multimap_reduce_array_container(
+                np.prod,
+                lambda x, y: x.shape == y.shape,
+                ary, ary_roundtrip)
+
+        assert actx.to_numpy(
+                actx.np.linalg.norm(ary - ary_roundtrip)
+                ) < 1.0e-15
+
+    # {{{ complex to real
+
+    if isinstance(shapes, int | tuple):
+        shapes = [shapes]
+
+    ary = DOFArray(actx, tuple(actx.from_numpy(randn(shape, np.float64))
+        for shape in shapes))
+
+    template = DOFArray(actx, tuple(actx.from_numpy(randn(shape, np.complex128))
+        for shape in shapes))
+
+    flat = flatten(ary, actx)
+    ary_roundtrip = unflatten(template, flat, actx, strict=False)
+
+    assert actx.to_numpy(
+            actx.np.linalg.norm(ary - ary_roundtrip)
+            ) < 1.0e-15
+
+    # }}}
+
+
+def _checked_flatten(ary, actx, leaf_class=None):
+    from arraycontext import flat_size_and_dtype, flatten
+    result = flatten(ary, actx, leaf_class=leaf_class)
+
+    if leaf_class is None:
+        size, dtype = flat_size_and_dtype(ary)
+        assert result.shape == (size,)
+        assert result.dtype == dtype
+
+    return result
+
+
+def test_flatten_array_container_failure(actx_factory):
+    actx = actx_factory()
+
+    from arraycontext import unflatten
+    ary = _get_test_containers(actx, shapes=512)[0]
+    flat_ary = _checked_flatten(ary, actx)
+
+    if not isinstance(actx, NumpyArrayContext):
+        with pytest.raises(TypeError):
+            # cannot unflatten from a numpy array (except for numpy actx)
+            unflatten(ary, actx.to_numpy(flat_ary), actx)
+
+    with pytest.raises(ValueError):
+        # cannot unflatten non-flat arrays
+        unflatten(ary, flat_ary.reshape(2, -1), actx)
+
+    with pytest.raises(ValueError):
+        # cannot unflatten partially
+        unflatten(ary, flat_ary[:-1], actx)
+
+
+def test_flatten_with_leaf_class(actx_factory):
+    actx = actx_factory()
+
+    arys = _get_test_containers(actx, shapes=512)
+
+    flat = _checked_flatten(arys[0], actx, leaf_class=DOFArray)
+    assert isinstance(flat, actx.array_types)
+    assert flat.shape == (arys[0].size,)
+
+    flat = _checked_flatten(arys[1], actx, leaf_class=DOFArray)
+    assert isinstance(flat, np.ndarray) and flat.dtype == object
+    assert all(isinstance(entry, actx.array_types) for entry in flat)
+    assert all(entry.shape == (arys[0].size,) for entry in flat)
+
+    flat = _checked_flatten(arys[3], actx, leaf_class=DOFArray)
+    assert isinstance(flat, MyContainer)
+    assert isinstance(flat.mass, actx.array_types)
+    assert flat.mass.shape == (arys[3].mass.size,)
+    assert isinstance(flat.enthalpy, actx.array_types)
+    assert flat.enthalpy.shape == (arys[3].enthalpy.size,)
+    assert all(isinstance(entry, actx.array_types) for entry in flat.momentum)
+
+# }}}
+
+
+# {{{ test from_numpy and to_numpy
+
+def test_numpy_conversion(actx_factory):
+    actx = actx_factory()
+    rng = np.random.default_rng()
+
+    nelements = 42
+    ac = MyContainer(
+            name="test_numpy_conversion",
+            mass=rng.uniform(size=(nelements, nelements)),
+            momentum=make_obj_array([rng.uniform(size=nelements) for _ in range(3)]),
+            enthalpy=np.array(rng.uniform()),
+            )
+
+    ac_actx = actx.from_numpy(ac)
+    ac_roundtrip = actx.to_numpy(ac_actx)
+
+    assert np.allclose(ac.mass, ac_roundtrip.mass)
+    assert np.allclose(ac.momentum[0], ac_roundtrip.momentum[0])
+
+    if not isinstance(actx, NumpyArrayContext):
+        from dataclasses import replace
+        ac_with_cl = replace(ac, enthalpy=ac_actx.mass)
+        with pytest.raises(TypeError):
+            actx.from_numpy(ac_with_cl)
+
+        with pytest.raises(TypeError):
+            actx.from_numpy(ac_actx)
+
+        with pytest.raises(TypeError):
+            actx.to_numpy(ac)
+
+# }}}
+
+
+# {{{ test actx.np.linalg.norm
+
+@pytest.mark.parametrize("norm_ord", [2, np.inf])
+def test_norm_complex(actx_factory, norm_ord):
+    actx = actx_factory()
+    a = randn(2000, np.complex128)
+
+    norm_a_ref = np.linalg.norm(a, norm_ord)
+    norm_a = actx.np.linalg.norm(actx.from_numpy(a), norm_ord)
+
+    norm_a = actx.to_numpy(norm_a)
+
+    assert abs(norm_a_ref - norm_a)/norm_a < 1e-13
+
+
+@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
+def test_norm_ord_none(actx_factory, ndim):
+    actx = actx_factory()
+
+    from numpy.random import default_rng
+
+    rng = default_rng()
+    shape = tuple(rng.integers(2, 7, ndim))
+
+    a = rng.random(shape)
+
+    norm_a_ref = np.linalg.norm(a, ord=None)
+    norm_a = actx.np.linalg.norm(actx.from_numpy(a), ord=None)
+
+    np.testing.assert_allclose(actx.to_numpy(norm_a), norm_a_ref)
+
+# }}}
+
+
+# {{{ test_actx_compile helpers
+
+def scale_and_orthogonalize(alpha, vel):
+    from arraycontext import rec_map_array_container
+    actx = vel.array_context
+    scaled_vel = rec_map_array_container(lambda x: alpha * x,
+                                         vel)
+    return Velocity2D(-scaled_vel.v, scaled_vel.u, actx)
+
+
+def test_actx_compile(actx_factory):
+    actx = actx_factory()
+    rng = np.random.default_rng()
+
+    compiled_rhs = actx.compile(scale_and_orthogonalize)
+
+    v_x = rng.uniform(size=10)
+    v_y = rng.uniform(size=10)
+
+    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
+
+    scaled_speed = compiled_rhs(np.float64(3.14), vel)
+
+    result = actx.to_numpy(scaled_speed)
+    np.testing.assert_allclose(result.u, -3.14*v_y)
+    np.testing.assert_allclose(result.v, 3.14*v_x)
+
+
+def test_actx_compile_python_scalar(actx_factory):
+    actx = actx_factory()
+    rng = np.random.default_rng()
+
+    compiled_rhs = actx.compile(scale_and_orthogonalize)
+
+    v_x = rng.uniform(size=10)
+    v_y = rng.uniform(size=10)
+
+    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
+
+    scaled_speed = compiled_rhs(3.14, vel)
+
+    result = actx.to_numpy(scaled_speed)
+    np.testing.assert_allclose(result.u, -3.14*v_y)
+    np.testing.assert_allclose(result.v, 3.14*v_x)
+
+
+def test_actx_compile_kwargs(actx_factory):
+    actx = actx_factory()
+    rng = np.random.default_rng()
+
+    compiled_rhs = actx.compile(scale_and_orthogonalize)
+
+    v_x = rng.uniform(size=10)
+    v_y = rng.uniform(size=10)
+
+    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
+
+    scaled_speed = compiled_rhs(3.14, vel=vel)
+
+    result = actx.to_numpy(scaled_speed)
+    np.testing.assert_allclose(result.u, -3.14*v_y)
+    np.testing.assert_allclose(result.v, 3.14*v_x)
+
+
+def test_actx_compile_with_tuple_output_keys(actx_factory):
+    # arraycontext.git<=3c9aee68 would fail due to a bug in output
+    # key stringification logic.
+    actx = actx_factory()
+    rng = np.random.default_rng()
+
+    def my_rhs(scale, vel):
+        result = np.empty((1, 1), dtype=object)
+        result[0, 0] = scale_and_orthogonalize(scale, vel)
+        return result
+
+    compiled_rhs = actx.compile(my_rhs)
+
+    v_x = rng.uniform(size=10)
+    v_y = rng.uniform(size=10)
+
+    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
+
+    scaled_speed = compiled_rhs(3.14, vel=vel)
+
+    result = actx.to_numpy(scaled_speed)[0, 0]
+    np.testing.assert_allclose(result.u, -3.14*v_y)
+    np.testing.assert_allclose(result.v, 3.14*v_x)
+
+# }}}
+
+
+# {{{ test_container_equality
+
+def test_container_equality(actx_factory):
+    actx = actx_factory()
+
+    ary_dof, _, _, _dc_of_dofs, bcast_dc_of_dofs = \
+            _get_test_containers(actx)
+    _, _, _, _dc_of_dofs_2, bcast_dc_of_dofs_2 = \
+            _get_test_containers(actx)
+
+    # MyContainer sets eq_comparison to False, so equality comparison should
+    # not succeed.
+    dc = MyContainer(name="yoink", mass=ary_dof, momentum=None, enthalpy=None)
+    dc2 = MyContainer(name="yoink", mass=ary_dof, momentum=None, enthalpy=None)
+    assert dc != dc2
+
+    assert isinstance(actx.np.equal(bcast_dc_of_dofs, bcast_dc_of_dofs_2),
+                      MyContainerDOFBcast)
+
+# }}}
+
+
+# {{{ test_leaf_array_type_broadcasting
+
+def test_no_leaf_array_type_broadcasting(actx_factory):
+    from testlib import Foo
+    # 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)
+    bar = foo + 4
+
+    bcast = partial(BcastUntilActxArray, actx)
+
+    actx_ary = actx.from_numpy(4*np.ones((3, )))
+    with pytest.raises(TypeError):
+        foo + actx_ary
+
+    baz = foo + bcast(actx_ary)
+    qux = bcast(actx_ary) + foo
+
+    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                               actx.to_numpy(baz.u[0]))
+
+    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                               actx.to_numpy(qux.u[0]))
+
+    baz = foo + bcast(actx_ary)
+    qux = bcast(actx_ary) + foo
+
+    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                               actx.to_numpy(baz.u[0]))
+
+    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                               actx.to_numpy(qux.u[0]))
+
+    mc = MyContainer(
+         name="hi",
+         mass=dof_ary,
+         momentum=make_obj_array([dof_ary, dof_ary]),
+         enthalpy=dof_ary)
+
+    with pytest.raises(TypeError):
+        mc_op = mc + actx_ary
+
+    mom_op = mc.momentum + bcast(actx_ary)
+    np.testing.assert_allclose(45, actx.to_numpy(mom_op[0][0]))
+
+    mc_op = mc + bcast(actx_ary)
+    np.testing.assert_allclose(45, actx.to_numpy(mc_op.mass[0]))
+    np.testing.assert_allclose(45, actx.to_numpy(mc_op.momentum[1][0]))
+
+    mom_op = mc.momentum + bcast(actx_ary)
+    np.testing.assert_allclose(45, actx.to_numpy(mom_op[0][0]))
+
+    mc_op = mc + bcast(actx_ary)
+    np.testing.assert_allclose(45, actx.to_numpy(mc_op.mass[0]))
+    np.testing.assert_allclose(45, actx.to_numpy(mc_op.momentum[1][0]))
+
+    def _actx_allows_scalar_broadcast(actx):
+        if not isinstance(actx, PyOpenCLArrayContext):
+            return True
+        else:
+            import pyopencl as cl
+
+            # See https://github.com/inducer/pyopencl/issues/498
+            return cl.version.VERSION > (2021, 2, 5)
+
+    if _actx_allows_scalar_broadcast(actx):
+        with pytest.raises(TypeError):
+            foo + actx.from_numpy(np.array(4))
+
+        quuz = bcast(actx.from_numpy(np.array(4))) + foo
+        quux = foo + bcast(actx.from_numpy(np.array(4)))
+
+        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                                   actx.to_numpy(quux.u[0]))
+
+        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                                   actx.to_numpy(quuz.u[0]))
+
+        quuz = bcast(actx.from_numpy(np.array(4))) + foo
+        quux = foo + bcast(actx.from_numpy(np.array(4)))
+
+        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                                   actx.to_numpy(quux.u[0]))
+
+        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
+                                   actx.to_numpy(quuz.u[0]))
+# }}}
+
+
+# {{{ test outer product
+
+def test_outer(actx_factory):
+    actx = actx_factory()
+
+    a_dof, a_ary_of_dofs, _, _, a_bcast_dc_of_dofs = _get_test_containers(actx)
+
+    b_dof = a_dof + 1
+    b_ary_of_dofs = a_ary_of_dofs + 1
+    b_bcast_dc_of_dofs = a_bcast_dc_of_dofs + 1
+
+    bcast = partial(BcastUntilActxArray, actx)
+
+    from arraycontext import outer
+
+    def equal(a, b):
+        return actx.to_numpy(actx.np.array_equal(a, b))
+
+    # Two scalars
+    assert equal(outer(a_dof, b_dof), a_dof*b_dof)
+
+    # Scalar and vector
+    assert equal(outer(a_dof, b_ary_of_dofs), a_dof*b_ary_of_dofs)
+
+    # Vector and scalar
+    assert equal(outer(a_ary_of_dofs, b_dof), a_ary_of_dofs*b_dof)
+
+    # Two vectors
+    assert equal(
+        outer(a_ary_of_dofs, b_ary_of_dofs),
+        np.outer(a_ary_of_dofs, b_ary_of_dofs))
+
+    # Scalar and array container
+    assert equal(
+        outer(a_dof, b_bcast_dc_of_dofs),
+        a_dof*b_bcast_dc_of_dofs)
+
+    # Array container and scalar
+    assert equal(
+        outer(a_bcast_dc_of_dofs, b_dof),
+        a_bcast_dc_of_dofs*b_dof)
+
+    # Vector and array container
+    assert equal(
+        outer(a_ary_of_dofs, b_bcast_dc_of_dofs),
+        make_obj_array([a_i*b_bcast_dc_of_dofs for a_i in a_ary_of_dofs]))
+
+    # Array container and vector
+    assert equal(
+        outer(a_bcast_dc_of_dofs, b_ary_of_dofs),
+        MyContainerDOFBcast(
+            name="container",
+            mass=a_bcast_dc_of_dofs.mass*b_ary_of_dofs,
+            momentum=np.outer(a_bcast_dc_of_dofs.momentum, b_ary_of_dofs),
+            enthalpy=a_bcast_dc_of_dofs.enthalpy*b_ary_of_dofs))
+
+    # Two array containers
+    assert equal(
+        outer(a_bcast_dc_of_dofs, b_bcast_dc_of_dofs),
+        MyContainerDOFBcast(
+            name="container",
+            mass=a_bcast_dc_of_dofs.mass*b_bcast_dc_of_dofs.mass,
+            momentum=np.outer(
+                a_bcast_dc_of_dofs.momentum,
+                b_bcast_dc_of_dofs.momentum),
+            enthalpy=a_bcast_dc_of_dofs.enthalpy*b_bcast_dc_of_dofs.enthalpy))
+
+    # Array context scalars
+    two = actx.from_numpy(np.array(2))
+    assert equal(
+        outer(bcast(two), b_bcast_dc_of_dofs),
+        bcast(two)*b_bcast_dc_of_dofs)
+    assert equal(
+        outer(a_bcast_dc_of_dofs, bcast(two)),
+        a_bcast_dc_of_dofs*bcast(two))
+
+# }}}
+
+
+# {{{ test_array_container_with_numpy
+
+@with_container_arithmetic(bcasts_across_obj_array=True, rel_comparison=True)
+@dataclass_array_container
+@dataclass(frozen=True)
+class ArrayContainerWithNumpy:
+    u: np.ndarray
+    v: DOFArray
+
+    __array_ufunc__ = None
+
+
+def test_array_container_with_numpy(actx_factory):
+    actx = actx_factory()
+
+    mystate = ArrayContainerWithNumpy(
+            u=np.zeros(10),
+            v=DOFArray(actx, (actx.from_numpy(np.zeros(42)),)),
+            )
+
+    from arraycontext import rec_map_array_container
+    rec_map_array_container(lambda x: x, mystate)
+
+
+# }}}
+
+
+# {{{ test_actx_compile_on_pure_array_return
+
+def test_actx_compile_on_pure_array_return(actx_factory):
+    def _twice(x):
+        return 2 * x
+
+    actx = actx_factory()
+    ones = actx.thaw(actx.freeze(
+        actx.np.zeros(shape=(10, 4), dtype=np.float64) + 1
+        ))
+    np.testing.assert_allclose(actx.to_numpy(_twice(ones)),
+                               actx.to_numpy(actx.compile(_twice)(ones)))
+
+# }}}
+
+
+# {{{ test_taggable_cl_array_tags
+
+@dataclass(frozen=True)
+class MySampleTag(Tag):
+    pass
+
+
+def test_taggable_cl_array_tags(actx_factory):
+    actx = actx_factory()
+    if not isinstance(actx, PyOpenCLArrayContext):
+        pytest.skip(f"not relevant for '{type(actx).__name__}'")
+
+    import pyopencl.array as cl_array
+    ary = cl_array.to_device(actx.queue, np.zeros((32, 7)))
+
+    # {{{ check tags are set
+
+    from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
+    tagged_ary = to_tagged_cl_array(ary, axes=None,
+                                    tags=frozenset((MySampleTag(),)))
+
+    assert tagged_ary.base_data is ary.base_data
+    assert tagged_ary.tags == frozenset((MySampleTag(),))
+
+    # }}}
+
+    # {{{ check tags are appended
+
+    from arraycontext import ElementwiseMapKernelTag
+    tagged_ary = to_tagged_cl_array(tagged_ary, axes=None,
+                                    tags=frozenset((ElementwiseMapKernelTag(),)))
+
+    assert tagged_ary.base_data is ary.base_data
+    assert tagged_ary.tags == frozenset(
+        (MySampleTag(), ElementwiseMapKernelTag())
+    )
+
+    # }}}
+
+    # {{{ test copied tags
+
+    copy_tagged_ary = tagged_ary.copy()
+
+    assert copy_tagged_ary.tags == tagged_ary.tags
+    assert copy_tagged_ary.axes == tagged_ary.axes
+    assert copy_tagged_ary.base_data != tagged_ary.base_data
+
+    # }}}
+
+# }}}
+
+
+def test_to_numpy_on_frozen_arrays(actx_factory):
+    # See https://github.com/inducer/arraycontext/issues/159
+    actx = actx_factory()
+    u = actx.freeze(actx.np.zeros(10, dtype="float64")+1)
+    np.testing.assert_allclose(actx.to_numpy(u), 1)
+    np.testing.assert_allclose(actx.to_numpy(u), 1)
+
+
+def test_tagging(actx_factory):
+    actx = actx_factory()
+
+    if isinstance(actx, NumpyArrayContext | EagerJAXArrayContext):
+        pytest.skip(f"{type(actx)} has no tagging support")
+
+    from pytools.tag import Tag
+
+    class ExampleTag(Tag):
+        pass
+
+    ary = tag_axes(actx, {0: ExampleTag()},
+            actx.tag(
+                ExampleTag(),
+                actx.np.zeros((20, 20), dtype=np.float64)))
+
+    assert ary.tags_of_type(ExampleTag)
+    assert ary.axes[0].tags_of_type(ExampleTag)
+    assert not ary.axes[1].tags_of_type(ExampleTag)
+
+
+def test_compile_anonymous_function(actx_factory):
+    from functools import partial
+
+    # See https://github.com/inducer/grudge/issues/287
+    actx = actx_factory()
+
+    ones = actx.thaw(actx.freeze(
+        actx.np.zeros(shape=(10, 4), dtype=np.float64) + 1
+        ))
+
+    f = actx.compile(lambda x: 2*x+40)
+    np.testing.assert_allclose(actx.to_numpy(f(ones)), 42)
+
+    f = actx.compile(partial(lambda x: 2*x+40))
+    np.testing.assert_allclose(actx.to_numpy(f(ones)), 42)
+
+
+@pytest.mark.parametrize(
+        ("args", "kwargs"), [
+            ((1, 2, 10), {}),
+            ((1, 2, 10), {"endpoint": False}),
+            ((1, 2, 10), {"endpoint": True}),
+            ((2, -3, 20), {}),
+            ((1, 5j, 20), {"dtype": np.complex128}),
+            ((1, 5, 20), {"dtype": np.complex128}),
+            ((1, 5, 20), {"dtype": np.int32}),
+            ])
+def test_linspace(actx_factory, args, kwargs):
+    if "Jax" in actx_factory.__class__.__name__:
+        pytest.xfail("jax actx does not have arange")
+
+    actx = actx_factory()
+
+    actx_linspace = actx.to_numpy(actx.np.linspace(*args, **kwargs))
+    np_linspace = np.linspace(*args, **kwargs)
+
+    assert np.allclose(actx_linspace, np_linspace)
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1:

test/test_pytato_arraycontext.py

+""" PytatoArrayContext specific tests"""
+from __future__ import annotations
+
+
+__copyright__ = "Copyright (C) 2021 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.
+"""
+
+import logging
+
+import numpy as np
+import pytest
+
+from pytools.tag import Tag
+
+from arraycontext import (
+    PytatoPyOpenCLArrayContext,
+    pytest_generate_tests_for_array_contexts,
+)
+from arraycontext.pytest import _PytestPytatoPyOpenCLArrayContextFactory
+
+
+logger = logging.getLogger(__name__)
+
+
+# {{{ pytato-array context fixture
+
+class _PytatoPyOpenCLArrayContextForTests(PytatoPyOpenCLArrayContext):
+    """Like :class:`PytatoPyOpenCLArrayContext`, but applies no program
+    transformations whatsoever. Only to be used for testing internal to
+    :mod:`arraycontext`.
+    """
+
+    def transform_loopy_program(self, t_unit):
+        return t_unit
+
+
+class _PytatoPyOpenCLArrayContextForTestsFactory(
+        _PytestPytatoPyOpenCLArrayContextFactory):
+    actx_class = _PytatoPyOpenCLArrayContextForTests
+
+
+pytest_generate_tests = pytest_generate_tests_for_array_contexts([
+    _PytatoPyOpenCLArrayContextForTestsFactory,
+    ])
+
+# }}}
+
+
+# {{{ dummy tag types
+
+class FooTag(Tag):
+    """
+    Foo
+    """
+
+
+class BarTag(Tag):
+    """
+    Bar
+    """
+
+
+class BazTag(Tag):
+    """
+    Baz
+    """
+
+# }}}
+
+
+def test_tags_preserved_after_freeze(actx_factory):
+    actx = actx_factory()
+
+    from arraycontext.impl.pytato import _BasePytatoArrayContext
+    if not isinstance(actx, _BasePytatoArrayContext):
+        pytest.skip("only pytato-based array context are supported")
+
+    from numpy.random import default_rng
+    rng = default_rng()
+
+    foo = actx.thaw(actx.freeze(
+        actx.from_numpy(rng.random((10, 4)))
+        .tagged(FooTag())
+        .with_tagged_axis(0, BarTag())
+        .with_tagged_axis(1, BazTag())
+        ))
+
+    assert foo.tags_of_type(FooTag)
+    assert foo.axes[0].tags_of_type(BarTag)
+    assert foo.axes[1].tags_of_type(BazTag)
+
+
+def test_arg_size_limit(actx_factory):
+    ran_callback = False
+
+    def my_ctc(what, stage, ir):
+        if stage == "final":
+            assert ir.target.limit_arg_size_nbytes == 42
+            nonlocal ran_callback
+            ran_callback = True
+
+    def twice(x):
+        return 2 * x
+
+    actx = _PytatoPyOpenCLArrayContextForTests(
+        actx_factory().queue, compile_trace_callback=my_ctc, _force_svm_arg_limit=42)
+
+    f = actx.compile(twice)
+    f(99)
+
+    assert ran_callback
+
+
+@pytest.mark.parametrize("pass_allocator", ["auto_none", "auto_true", "auto_false",
+                                            "pass_buffer", "pass_svm",
+                                            "pass_buffer_pool", "pass_svm_pool"])
+def test_pytato_actx_allocator(actx_factory, pass_allocator):
+    base_actx = actx_factory()
+    alloc = None
+    use_memory_pool = None
+
+    if pass_allocator == "auto_none":
+        pass
+    elif pass_allocator == "auto_true":
+        use_memory_pool = True
+    elif pass_allocator == "auto_false":
+        use_memory_pool = False
+    elif pass_allocator == "pass_buffer":
+        from pyopencl.tools import ImmediateAllocator
+        alloc = ImmediateAllocator(base_actx.queue)
+    elif pass_allocator == "pass_svm":
+        from pyopencl.characterize import has_coarse_grain_buffer_svm
+        if not has_coarse_grain_buffer_svm(base_actx.queue.device):
+            pytest.skip("need SVM support for this test")
+        from pyopencl.tools import SVMAllocator
+        alloc = SVMAllocator(base_actx.queue.context, queue=base_actx.queue)
+    elif pass_allocator == "pass_buffer_pool":
+        from pyopencl.tools import ImmediateAllocator, MemoryPool
+        alloc = MemoryPool(ImmediateAllocator(base_actx.queue))
+    elif pass_allocator == "pass_svm_pool":
+        from pyopencl.characterize import has_coarse_grain_buffer_svm
+        if not has_coarse_grain_buffer_svm(base_actx.queue.device):
+            pytest.skip("need SVM support for this test")
+        from pyopencl.tools import SVMAllocator, SVMPool
+        alloc = SVMPool(SVMAllocator(base_actx.queue.context, queue=base_actx.queue))
+    else:
+        raise ValueError(f"unknown option {pass_allocator}")
+
+    actx = _PytatoPyOpenCLArrayContextForTests(base_actx.queue, allocator=alloc,
+                                               use_memory_pool=use_memory_pool)
+
+    def twice(x):
+        return 2 * x
+
+    f = actx.compile(twice)
+    res = actx.to_numpy(f(99))
+
+    assert res == 198
+
+    # Also test a case in which SVM is not available
+    if pass_allocator in ["auto_none", "auto_true", "auto_false"]:
+        from unittest.mock import patch
+
+        with patch("pyopencl.characterize.has_coarse_grain_buffer_svm",
+                    return_value=False):
+            actx = _PytatoPyOpenCLArrayContextForTests(base_actx.queue,
+                        allocator=alloc, use_memory_pool=use_memory_pool)
+
+            from pyopencl.tools import ImmediateAllocator, MemoryPool
+            assert isinstance(actx.allocator,
+                              MemoryPool if use_memory_pool else ImmediateAllocator)
+
+            f = actx.compile(twice)
+            res = actx.to_numpy(f(99))
+
+            assert res == 198
+
+
+def test_transfer(actx_factory):
+    import numpy as np
+
+    import pytato as pt
+    actx = actx_factory()
+
+    # {{{ simple tests
+
+    a = actx.from_numpy(np.array([0, 1, 2, 3])).tagged(FooTag())
+
+    from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
+    assert isinstance(a.data, TaggableCLArray)
+
+    from arraycontext.impl.pytato.utils import transfer_from_numpy, transfer_to_numpy
+
+    ah = transfer_to_numpy(a, actx)
+    assert ah != a
+    assert a.tags == ah.tags
+    assert a.non_equality_tags == ah.non_equality_tags
+    assert isinstance(ah.data, np.ndarray)
+
+    with pytest.raises(ValueError):
+        _ahh = transfer_to_numpy(ah, actx)
+
+    ad = transfer_from_numpy(ah, actx)
+    assert isinstance(ad.data, TaggableCLArray)
+    assert ad != ah
+    assert ad != a  # copied DataWrappers compare unequal
+    assert ad.tags == ah.tags
+    assert ad.non_equality_tags == ah.non_equality_tags
+    assert np.array_equal(a.data.get(), ad.data.get())
+
+    with pytest.raises(ValueError):
+        _add = transfer_from_numpy(ad, actx)
+
+    # }}}
+
+    # {{{ test with DictOfNamedArrays
+
+    dag = pt.make_dict_of_named_arrays({
+        "a_expr": a + 2
+        })
+
+    dagh = transfer_to_numpy(dag, actx)
+    assert dagh != dag
+    assert isinstance(dagh["a_expr"].expr.bindings["_in0"].data, np.ndarray)
+
+    daghd = transfer_from_numpy(dagh, actx)
+    assert isinstance(daghd["a_expr"].expr.bindings["_in0"].data, TaggableCLArray)
+
+    # }}}
+
+
+def test_pass_args_compiled_func(actx_factory):
+    import numpy as np
+
+    import loopy as lp
+    import pyopencl as cl
+    import pyopencl.array
+    import pytato as pt
+
+    def twice(x, y, a):
+        return 2 * x * y * a
+
+    actx = _PytatoPyOpenCLArrayContextForTests(actx_factory().queue)
+
+    dev_scalar = pt.make_data_wrapper(cl.array.to_device(actx.queue, np.float64(23)))
+
+    f = actx.compile(twice)
+
+    assert actx.to_numpy(f(99.0, np.float64(2.0), dev_scalar)) == 2*23*99*2
+
+    compiled_func, = f.program_cache.values()
+    ep = compiled_func.pytato_program.program.t_unit.default_entrypoint
+
+    assert isinstance(ep.arg_dict["_actx_in_0"], lp.ValueArg)
+    assert isinstance(ep.arg_dict["_actx_in_1"], lp.ValueArg)
+    assert isinstance(ep.arg_dict["_actx_in_2"], lp.ArrayArg)
+
+
+def test_profiling_actx():
+    import pyopencl as cl
+    cl_ctx = cl.create_some_context()
+    queue = cl.CommandQueue(cl_ctx,
+                            properties=cl.command_queue_properties.PROFILING_ENABLE)
+
+    actx = PytatoPyOpenCLArrayContext(queue, profile_kernels=True)
+
+    def twice(x):
+        return 2 * x
+
+    # {{{ Compiled test
+
+    f = actx.compile(twice)
+
+    assert len(actx._profile_events) == 0
+
+    for _ in range(10):
+        assert actx.to_numpy(f(99)) == 198
+
+    assert len(actx._profile_events) == 10
+    actx._wait_and_transfer_profile_events()
+    assert len(actx._profile_events) == 0
+    assert len(actx._profile_results) == 1
+    assert len(actx._profile_results["twice"]) == 10
+
+    from arraycontext.impl.pytato.utils import tabulate_profiling_data
+
+    print(tabulate_profiling_data(actx))
+    assert len(actx._profile_results) == 0
+
+    # }}}
+
+    # {{{ Uncompiled/frozen test
+
+    assert len(actx._profile_events) == 0
+
+    for _ in range(10):
+        assert np.all(actx.to_numpy(twice(actx.from_numpy(np.array([99, 99])))) == 198)
+
+    assert len(actx._profile_events) == 10
+    actx._wait_and_transfer_profile_events()
+    assert len(actx._profile_events) == 0
+    assert len(actx._profile_results) == 1
+    assert len(actx._profile_results["frozen_result"]) == 10
+
+    print(tabulate_profiling_data(actx))
+
+    assert len(actx._profile_results) == 0
+
+    # }}}
+
+    # {{{ test disabling profiling
+
+    actx._enable_profiling(False)
+
+    assert len(actx._profile_events) == 0
+
+    for _ in range(10):
+        assert actx.to_numpy(f(99)) == 198
+
+    assert len(actx._profile_events) == 0
+    assert len(actx._profile_results) == 0
+
+    # }}}
+
+    # {{{ test enabling profiling
+
+    actx._enable_profiling(True)
+
+    assert len(actx._profile_events) == 0
+
+    for _ in range(10):
+        assert actx.to_numpy(f(99)) == 198
+
+    assert len(actx._profile_events) == 10
+    actx._wait_and_transfer_profile_events()
+    assert len(actx._profile_events) == 0
+    assert len(actx._profile_results) == 1
+
+    # }}}
+
+    queue2 = cl.CommandQueue(cl_ctx)
+
+    with pytest.raises(RuntimeError):
+        PytatoPyOpenCLArrayContext(queue2, profile_kernels=True)
+
+    actx2 = PytatoPyOpenCLArrayContext(queue2)
+
+    with pytest.raises(RuntimeError):
+        actx2._enable_profiling(True)
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1:
+        exec(sys.argv[1])
+    else:
+        pytest.main([__file__])
+
+# vim: fdm=marker

test/test_utils.py

+"""Testing for internal  utilities."""
+
+# Do not add
+# from __future__ import annotations
+# to allow the non-string annotations below to work.
+
+
+__copyright__ = "Copyright (C) 2021 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.
+"""
+import logging
+from typing import Optional, Tuple, cast  # noqa: UP035
+
+import numpy as np
+import pytest
+
+
+logger = logging.getLogger(__name__)
+
+
+# {{{ test_pt_actx_key_stringification_uniqueness
+
+def test_pt_actx_key_stringification_uniqueness():
+    from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
+
+    assert (_ary_container_key_stringifier(((3, 2), 3))
+            != _ary_container_key_stringifier((3, (2, 3))))
+
+    assert (_ary_container_key_stringifier(("tup", 3, "endtup"))
+            != _ary_container_key_stringifier(((3,),)))
+
+# }}}
+
+
+# {{{ test_dataclass_array_container
+
+def test_dataclass_array_container() -> None:
+    from dataclasses import dataclass, field
+
+    from arraycontext import Array, dataclass_array_container
+
+    # {{{ optional fields
+
+    @dataclass
+    class ArrayContainerWithOptional:
+        x: np.ndarray
+        # Deliberately left as Optional to test compatibility.
+        y: Optional[np.ndarray]  # noqa: UP045
+
+    with pytest.raises(TypeError, match="Field 'y' union contains non-array"):
+        # NOTE: cannot have wrapped annotations (here by `Optional`)
+        dataclass_array_container(ArrayContainerWithOptional)
+
+    # }}}
+
+    # {{{ type annotations
+
+    @dataclass
+    class ArrayContainerWithTuple:
+        x: Array
+        # Deliberately left as Tuple to test compatibility.
+        y: Tuple[Array, Array]  # noqa: UP006
+
+    with pytest.raises(TypeError, match="Typing annotation not supported on field 'y'"):
+        dataclass_array_container(ArrayContainerWithTuple)
+
+    @dataclass
+    class ArrayContainerWithTupleAlt:
+        x: Array
+        y: tuple[Array, Array]
+
+    with pytest.raises(TypeError, match="Typing annotation not supported on field 'y'"):
+        dataclass_array_container(ArrayContainerWithTupleAlt)
+
+    # }}}
+
+    # {{{ field(init=False)
+
+    @dataclass
+    class ArrayContainerWithInitFalse:
+        x: np.ndarray
+        y: np.ndarray = field(default_factory=lambda: np.zeros(42),
+                              init=False, repr=False)
+
+    with pytest.raises(ValueError, match="Field with 'init=False' not allowed"):
+        # NOTE: init=False fields are not allowed
+        dataclass_array_container(ArrayContainerWithInitFalse)
+
+    # }}}
+
+    # {{{ device arrays
+
+    @dataclass
+    class ArrayContainerWithArray:
+        x: Array
+        y: Array
+
+    dataclass_array_container(ArrayContainerWithArray)
+
+    # }}}
+
+# }}}
+
+
+# {{{ test_dataclass_container_unions
+
+def test_dataclass_container_unions() -> None:
+    from dataclasses import dataclass
+
+    from arraycontext import Array, dataclass_array_container
+
+    # {{{ union fields
+
+    @dataclass
+    class ArrayContainerWithUnion:
+        x: np.ndarray
+        y: np.ndarray | Array
+
+    dataclass_array_container(ArrayContainerWithUnion)
+
+    @dataclass
+    class ArrayContainerWithUnionAlt:
+        x: np.ndarray
+        y: np.ndarray | Array
+
+    dataclass_array_container(ArrayContainerWithUnionAlt)
+
+    # }}}
+
+    # {{{ non-container union
+
+    @dataclass
+    class ArrayContainerWithWrongUnion:
+        x: np.ndarray
+        y: np.ndarray | float
+
+    with pytest.raises(TypeError, match="Field 'y' union contains non-array container"):
+        # NOTE: float is not an ArrayContainer, so y should fail
+        dataclass_array_container(ArrayContainerWithWrongUnion)
+
+    # }}}
+
+    # {{{ optional union
+
+    @dataclass
+    class ArrayContainerWithOptionalUnion:
+        x: np.ndarray
+        y: np.ndarray | None
+
+    with pytest.raises(TypeError, match="Field 'y' union contains non-array container"):
+        # NOTE: None is not an ArrayContainer, so y should fail
+        dataclass_array_container(ArrayContainerWithWrongUnion)
+
+    # }}}
+
+
+# }}}
+
+
+# {{{ test_stringify_array_container_tree
+
+
+def test_stringify_array_container_tree() -> None:
+    from dataclasses import dataclass
+
+    from arraycontext import (
+        Array,
+        dataclass_array_container,
+        stringify_array_container_tree,
+    )
+
+    @dataclass_array_container
+    @dataclass(frozen=True)
+    class ArrayWrapper:
+        ary: Array
+
+    @dataclass_array_container
+    @dataclass(frozen=True)
+    class SomeContainer:
+        points: Array
+        radius: float
+        centers: ArrayWrapper
+
+    @dataclass_array_container
+    @dataclass(frozen=True)
+    class SomeOtherContainer:
+        disk: SomeContainer
+        circle: SomeContainer
+        has_disk: bool
+        norm_type: str
+        extent: float
+
+    rng = np.random.default_rng(seed=42)
+    a = ArrayWrapper(ary=cast(Array, rng.random(10)))
+    d = SomeContainer(
+              points=cast(Array, rng.random((2, 10))),
+              radius=rng.random(),
+              centers=a)
+    c = SomeContainer(
+              points=cast(Array, rng.random((2, 10))),
+              radius=rng.random(),
+              centers=a)
+    ary = SomeOtherContainer(
+        disk=d, circle=c,
+        has_disk=True,
+        norm_type="l2",
+        extent=1)
+
+    logger.info("\n%s", stringify_array_container_tree(ary))
+
+# }}}
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1:
+        exec(sys.argv[1])
+    else:
+        pytest.main([__file__])
+
+# vim: fdm=marker

test/testlib.py

+from __future__ import annotations
+
+
+__copyright__ = "Copyright (C) 2020-21 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 dataclasses import dataclass
+
+import numpy as np
+
+from arraycontext import (
+    ArrayContainer,
+    ArrayContext,
+    dataclass_array_container,
+    deserialize_container,
+    serialize_container,
+    with_array_context,
+    with_container_arithmetic,
+)
+
+
+# Containers live here, because in order for get_annotations to work, they must
+# live somewhere importable.
+# See https://docs.python.org/3.12/library/inspect.html#inspect.get_annotations
+
+
+# {{{ stand-in DOFArray implementation
+
+@with_container_arithmetic(
+        bcasts_across_obj_array=True,
+        bitwise=True,
+        rel_comparison=True,
+        _cls_has_array_context_attr=True,
+        _bcast_actx_array_type=False)
+class DOFArray:
+    def __init__(self, actx, data):
+        if not (actx is None or isinstance(actx, ArrayContext)):
+            raise TypeError("actx must be of type ArrayContext")
+
+        if not isinstance(data, tuple):
+            raise TypeError("'data' argument must be a tuple")
+
+        self.array_context = actx
+        self.data = data
+
+    # prevent numpy broadcasting
+    __array_ufunc__ = None
+
+    def __bool__(self):
+        if len(self) == 1 and self.data[0].size == 1:
+            return bool(self.data[0])
+
+        raise ValueError(
+                "The truth value of an array with more than one element is "
+                "ambiguous. Use actx.np.any(x) or actx.np.all(x)")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        return self.data[i]
+
+    def __repr__(self):
+        return f"DOFArray({self.data!r})"
+
+    @classmethod
+    def _serialize_init_arrays_code(cls, instance_name):
+        return {"_":
+                (f"{instance_name}_i", f"{instance_name}")}
+
+    @classmethod
+    def _deserialize_init_arrays_code(cls, template_instance_name, args):
+        (_, arg), = args.items()
+        # Why tuple([...])? https://stackoverflow.com/a/48592299
+        return (f"{template_instance_name}.array_context, tuple([{arg}])")
+
+    @property
+    def size(self):
+        return sum(ary.size for ary in self.data)
+
+    @property
+    def real(self):
+        return DOFArray(self.array_context, tuple(subary.real for subary in self))
+
+    @property
+    def imag(self):
+        return DOFArray(self.array_context, tuple(subary.imag for subary in self))
+
+
+@serialize_container.register(DOFArray)
+def _serialize_dof_container(ary: DOFArray):
+    return list(enumerate(ary.data))
+
+
+@deserialize_container.register(DOFArray)
+# https://github.com/python/mypy/issues/13040
+def _deserialize_dof_container(  # type: ignore[misc]
+        template, iterable):
+    def _raise_index_inconsistency(i, stream_i):
+        raise ValueError(
+                "out-of-sequence indices supplied in DOFArray deserialization "
+                f"(expected {i}, received {stream_i})")
+
+    return type(template)(
+            template.array_context,
+            data=tuple(
+                v if i == stream_i else _raise_index_inconsistency(i, stream_i)
+                for i, (stream_i, v) in enumerate(iterable)))
+
+
+@with_array_context.register(DOFArray)
+# https://github.com/python/mypy/issues/13040
+def _with_actx_dofarray(ary: DOFArray, actx: ArrayContext) -> DOFArray:  # type: ignore[misc]
+    return type(ary)(actx, ary.data)
+
+# }}}
+
+
+# {{{ nested containers
+
+@with_container_arithmetic(bcasts_across_obj_array=False,
+        eq_comparison=False, rel_comparison=False,
+        _cls_has_array_context_attr=True,
+        _bcast_actx_array_type=False)
+@dataclass_array_container
+@dataclass(frozen=True)
+class MyContainer:
+    name: str
+    mass: DOFArray | np.ndarray
+    momentum: np.ndarray
+    enthalpy: DOFArray | np.ndarray
+
+    __array_ufunc__ = None
+
+    @property
+    def array_context(self):
+        if isinstance(self.mass, np.ndarray):
+            return next(iter(self.mass)).array_context
+        else:
+            return self.mass.array_context
+
+
+@with_container_arithmetic(
+        bcasts_across_obj_array=False,
+        container_types_bcast_across=(DOFArray, np.ndarray),
+        matmul=True,
+        rel_comparison=True,
+        _cls_has_array_context_attr=True,
+        _bcast_actx_array_type=False)
+@dataclass_array_container
+@dataclass(frozen=True)
+class MyContainerDOFBcast:
+    name: str
+    mass: DOFArray | np.ndarray
+    momentum: np.ndarray
+    enthalpy: DOFArray | np.ndarray
+
+    __array_ufunc__ = None
+
+    @property
+    def array_context(self):
+        if isinstance(self.mass, np.ndarray):
+            return next(iter(self.mass)).array_context
+        else:
+            return self.mass.array_context
+
+# }}}
+
+
+@with_container_arithmetic(
+    bcasts_across_obj_array=True,
+    rel_comparison=True,
+    _cls_has_array_context_attr=True,
+    _bcast_actx_array_type=False)
+@dataclass_array_container
+@dataclass(frozen=True)
+class Foo:
+    u: DOFArray
+
+    # prevent numpy arithmetic from taking precedence
+    __array_ufunc__ = None
+
+    @property
+    def array_context(self):
+        return self.u.array_context
+
+
+@with_container_arithmetic(bcasts_across_obj_array=True, rel_comparison=True)
+@dataclass_array_container
+@dataclass(frozen=True)
+class Velocity2D:
+    u: ArrayContainer
+    v: ArrayContainer
+    array_context: ArrayContext
+
+    __array_ufunc__ = None
+
+
+@with_array_context.register(Velocity2D)
+# https://github.com/python/mypy/issues/13040
+def _with_actx_velocity_2d(ary, actx):  # type: ignore[misc]
+    return type(ary)(ary.u, ary.v, actx)