run-pylint.sh

@@ -20,4 +20,4 @@ if [[ -f .pylintrc-local.yml ]]; then
     PYLINT_RUNNER_ARGS+=" --yaml-rcfile=.pylintrc-local.yml"
 fi
 
-python .run-pylint.py $PYLINT_RUNNER_ARGS $(basename $PWD) test/test_*.py examples "$@"
+PYTHONWARNINGS=ignore python .run-pylint.py $PYLINT_RUNNER_ARGS $(basename $PWD) test/test_*.py examples "$@"

setup.cfg

-[flake8]
-min_python_version = 3.6
-ignore = E126,E127,E128,E123,E226,E241,E242,E265,W503,E402
-max-line-length=85
-
-inline-quotes = "
-docstring-quotes = """
-multiline-quotes = """
-# enable-flake8-bugbear
-
-[mypy]
-python_version = 3.6
-warn_unused_ignores = True
-
-[mypy-islpy]
-ignore_missing_imports = True
-
-[mypy-loopy.*]
-ignore_missing_imports = True
-
-[mypy-numpy]
-ignore_missing_imports = True
-
-[mypy-meshmode.*]
-ignore_missing_imports = True
-
-[mypy-pymbolic.*]
-ignore_missing_imports = True
-
-[mypy-pyopencl.*]
-ignore_missing_imports = True

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 :: 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.7'",
-            "types-dataclasses",
-        ],
-        package_data={"arraycontext": ["py.typed"]},
-    )
-
-
-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,29 +23,37 @@ 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,
-        PyOpenCLArrayContext,
-        PytatoPyOpenCLArrayContext,
-        ArrayContainer,)
-from arraycontext import (  # noqa: F401
-        pytest_generate_tests_for_array_contexts,
-        )
-from arraycontext.pytest import (_PytestPyOpenCLArrayContextFactoryWithClass,
-                                 _PytestPytatoPyOpenCLArrayContextFactory)
+    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__)
 
 
@@ -86,6 +97,9 @@ pytest_generate_tests = pytest_generate_tests_for_array_contexts([
     _PyOpenCLArrayContextForTestsFactory,
     _PyOpenCLArrayContextWithHostScalarsForTestsFactory,
     _PytatoPyOpenCLArrayContextForTestsFactory,
+    _PytestEagerJaxArrayContextFactory,
+    _PytestPytatoJaxArrayContextFactory,
+    _PytestNumpyArrayContextFactory,
     ])
 
 
@@ -99,91 +113,38 @@ def _acf():
 # }}}
 
 
-# {{{ stand-in DOFArray implementation
-
-@with_container_arithmetic(
-        bcast_obj_array=True,
-        bcast_numpy_array=True,
-        bitwise=True,
-        rel_comparison=True,
-        _cls_has_array_context_attr=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")
-
-        self.array_context = actx
-        self.data = data
-
-    __array_priority__ = 10
-
-    def __len__(self):
-        return len(self.data)
-
-    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}")}
-
-    @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]))
-
-    @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 enumerate(ary.data)
-
-
-@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})")
-
-    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)))
+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]
 
-@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))
+    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)
 
-@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")
+    # 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)
 
-    return type(ary)(
-        actx,
-        tuple(actx.thaw(subary) for subary in ary.data))
+    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)
 
 
 # {{{ assert_close_to_numpy*
@@ -192,17 +153,24 @@ 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}")
-        return rng.standard_normal(shape, dtype) \
-            + 1j * rng.standard_normal(shape, dtype)
+        r = rng.standard_normal(ashape, dtype) \
+            + 1j * rng.standard_normal(ashape, dtype)
     elif dtype.kind == "f":
-        return rng.standard_normal(shape, dtype)
+        r = rng.standard_normal(ashape, dtype)
     elif dtype.kind == "i":
-        return rng.integers(0, 128, shape, dtype)
+        r = rng.integers(0, 512, ashape, dtype)
     else:
         raise TypeError(dtype.kind)
 
+    if shape == 0:
+        return np.array(r[0])
+
+    return r
+
 
 def assert_close_to_numpy(actx, op, args):
     assert np.allclose(
@@ -260,6 +228,9 @@ def assert_close_to_numpy_in_containers(actx, op, args):
             ("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),
@@ -274,6 +245,7 @@ def assert_close_to_numpy_in_containers(actx, op, args):
             ("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()
@@ -283,8 +255,23 @@ def test_array_context_np_workalike(actx_factory, sym_name, n_args, dtype):
     ndofs = 512
     args = [randn(ndofs, dtype) for i in range(n_args)]
 
-    assert_close_to_numpy_in_containers(
-            actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), 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)))
+
+        return func(*args_)
+
+    assert_close_to_numpy_in_containers(actx, evaluate, args)
+
+    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)
 
 
 @pytest.mark.parametrize(("sym_name", "n_args", "dtype"), [
@@ -301,38 +288,59 @@ 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)):
+        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}'")
+
 # }}}
 
 
 # {{{ array manipulations
 
 def test_actx_stack(actx_factory):
+    rng = np.random.default_rng()
+
     actx = actx_factory()
 
     ndofs = 5000
-    args = [np.random.randn(ndofs) for i in range(10)]
+    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)]
+    args = [rng.normal(size=ndofs) for i in range(10)]
 
     assert_close_to_numpy(
             actx, lambda _np, *_args: _np.concatenate(_args), args)
 
 
 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),
-                (np.random.randn(2, 3), new_shape))
+                (rng.normal(size=(2, 3)), new_shape))
 
 
 def test_actx_ravel(actx_factory):
@@ -351,6 +359,7 @@ def test_actx_ravel(actx_factory):
 # {{{ 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
@@ -362,8 +371,9 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
         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),
@@ -371,12 +381,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),
@@ -431,9 +441,9 @@ def test_dof_array_arithmetic_same_as_numpy(actx_factory):
                 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)
@@ -503,22 +513,72 @@ def test_dof_array_arithmetic_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()
 
-    ary = np.random.randn(3000)
+    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)
 
-    from numbers import Number
-
-    if isinstance(actx, PyOpenCLArrayContext) and (not actx._force_device_scalars):
-        assert isinstance(actx_red, Number)
-    else:
-        assert actx_red.shape == ()
+    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__}'")
+
+    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))
+
+
 # }}}
 
 
@@ -531,10 +591,10 @@ def test_reductions_same_as_numpy(actx_factory, op):
 ])
 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)
 
@@ -546,11 +606,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)
 
@@ -560,12 +620,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),
@@ -577,75 +637,109 @@ def test_array_context_einsum_array_tripleprod(actx_factory, spec):
 
 # {{{ array container classes for test
 
-@with_container_arithmetic(bcast_obj_array=False,
-        eq_comparison=False, rel_comparison=False)
-@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
-
-
-@with_container_arithmetic(
-        bcast_obj_array=False,
-        bcast_container_types=(DOFArray, np.ndarray),
-        matmul=True,
-        rel_comparison=True,)
-@dataclass_array_container
-@dataclass(frozen=True)
-class MyContainerDOFBcast:
-    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))
 
-    # pylint: disable=unexpected-keyword-arg, no-value-for-parameter
-    bcast_dataclass_of_dofs = MyContainerDOFBcast(
-            name="container",
-            mass=x,
-            momentum=make_obj_array([x, x]),
-            enthalpy=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
 
-    ary_dof = x
-    ary_of_dofs = make_obj_array([x, x, x])
-    mat_of_dofs = np.empty((3, 3), 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)
+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
+
+    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, bcast_dc_of_dofs = \
+    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=2.0e-14):
-        assert np.linalg.norm(actx.to_numpy(f(arg1) - arg2)) < atol
+        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
@@ -660,17 +754,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)
@@ -688,7 +795,6 @@ def test_container_arithmetic(actx_factory):
     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(
@@ -724,49 +830,52 @@ def test_container_arithmetic(actx_factory):
 
     grad_matvec_result = mock_gradient @ ary_of_dofs
     assert isinstance(grad_matvec_result.mass, DOFArray)
-    assert grad_matvec_result.momentum.shape == (3,)
+    assert grad_matvec_result.momentum.shape == ary_of_dofs.shape
 
-    assert actx.to_numpy(actx.np.linalg.norm(grad_matvec_result.mass
-                                             - 3*ary_of_dofs**2)) < 1e-8
+    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, bcast_dc_of_dofs = \
+    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(frozen_ary) is None
-        assert get_container_context_recursively(thawed_ary) is actx
+        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 = freeze(ary_dof)
+    ary_dof_frozen = actx.freeze(ary_dof)
     with pytest.raises(ValueError) as exc_info:
         ary_dof + ary_dof_frozen
 
     assert "frozen" in str(exc_info.value)
 
-    ary_dof_2 = thaw(freeze(ary_dof), actx2)
+    ary_dof_2 = actx2.thaw(actx.freeze(ary_dof))
 
     with pytest.raises(ValueError):
         ary_dof + ary_dof_2
@@ -788,35 +897,147 @@ 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=np.random.rand(42),
-            momentum=make_obj_array([np.random.rand(42) for _ in range(3)]),
-            enthalpy=np.random.rand(42),
+            mass=rng.uniform(size=(nelements, nelements)),
+            momentum=make_obj_array([rng.uniform(size=nelements) for _ in range(3)]),
+            enthalpy=np.array(rng.uniform()),
             )
 
-    from arraycontext import from_numpy, to_numpy
-    ac_actx = from_numpy(ac, actx)
-    ac_roundtrip = to_numpy(ac_actx, actx)
+    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])
 
-    from dataclasses import replace
-    ac_with_cl = replace(ac, enthalpy=ac_actx.mass)
-    with pytest.raises(TypeError):
-        from_numpy(ac_with_cl, actx)
+    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):
-        from_numpy(ac_actx, actx)
+        with pytest.raises(TypeError):
+            actx.from_numpy(ac_actx)
 
-    with pytest.raises(ValueError):
-        to_numpy(ac, actx)
+        with pytest.raises(TypeError):
+            actx.to_numpy(ac)
 
 # }}}
 
@@ -857,15 +1078,6 @@ def test_norm_ord_none(actx_factory, ndim):
 
 # {{{ test_actx_compile helpers
 
-@with_container_arithmetic(bcast_obj_array=True, rel_comparison=True)
-@dataclass_array_container
-@dataclass(frozen=True)
-class Velocity2D:
-    u: ArrayContainer
-    v: ArrayContainer
-    array_context: ArrayContext
-
-
 def scale_and_orthogonalize(alpha, vel):
     from arraycontext import rec_map_array_container
     actx = vel.array_context
@@ -875,19 +1087,80 @@ def scale_and_orthogonalize(alpha, vel):
 
 
 def test_actx_compile(actx_factory):
-    from arraycontext import (to_numpy, from_numpy)
     actx = actx_factory()
+    rng = np.random.default_rng()
 
     compiled_rhs = actx.compile(scale_and_orthogonalize)
 
-    v_x = np.random.rand(10)
-    v_y = np.random.rand(10)
+    v_x = rng.uniform(size=10)
+    v_y = rng.uniform(size=10)
 
-    vel = from_numpy(Velocity2D(v_x, v_y, actx), actx)
+    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
 
     scaled_speed = compiled_rhs(np.float64(3.14), vel)
 
-    result = to_numpy(scaled_speed, actx)
+    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)
 
@@ -899,9 +1172,9 @@ def test_actx_compile(actx_factory):
 def test_container_equality(actx_factory):
     actx = actx_factory()
 
-    ary_dof, _, _, dc_of_dofs, bcast_dc_of_dofs = \
+    ary_dof, _, _, _dc_of_dofs, bcast_dc_of_dofs = \
             _get_test_containers(actx)
-    _, _, _, dc_of_dofs_2, bcast_dc_of_dofs_2 = \
+    _, _, _, _dc_of_dofs_2, bcast_dc_of_dofs_2 = \
             _get_test_containers(actx)
 
     # MyContainer sets eq_comparison to False, so equality comparison should
@@ -910,64 +1183,263 @@ def test_container_equality(actx_factory):
     dc2 = MyContainer(name="yoink", mass=ary_dof, momentum=None, enthalpy=None)
     assert dc != dc2
 
-    assert isinstance(bcast_dc_of_dofs == bcast_dc_of_dofs_2, MyContainerDOFBcast)
+    assert isinstance(actx.np.equal(bcast_dc_of_dofs, bcast_dc_of_dofs_2),
+                      MyContainerDOFBcast)
 
 # }}}
 
 
-# {{{ test_leaf_array_type_broadcasting
+# {{{ test_no_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)
+
+    actx_ary = actx.from_numpy(4*np.ones((3, )))
+    with pytest.raises(TypeError):
+        foo + actx_ary
+
+    with pytest.raises(TypeError):
+        foo + actx.from_numpy(np.array(4))
+
+# }}}
+
 
-@with_container_arithmetic(
-    bcast_obj_array=True,
-    bcast_numpy_array=True,
-    rel_comparison=True,
-    _cls_has_array_context_attr=True)
+# {{{ 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
+
+    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))
+
+# }}}
+
+
+# {{{ test_array_container_with_numpy
+
+@with_container_arithmetic(bcasts_across_obj_array=True, rel_comparison=True)
 @dataclass_array_container
 @dataclass(frozen=True)
-class Foo:
-    u: DOFArray
+class ArrayContainerWithNumpy:
+    u: np.ndarray
+    v: DOFArray
 
-    @property
-    def array_context(self):
-        return self.u.array_context
+    __array_ufunc__ = None
 
 
-def test_leaf_array_type_broadcasting(actx_factory):
-    # test support for https://github.com/inducer/arraycontext/issues/49
+def test_array_container_with_numpy(actx_factory):
     actx = actx_factory()
 
-    foo = Foo(DOFArray(actx, (actx.zeros(3, dtype=np.float64) + 41, )))
-    bar = foo + 4
-    baz = foo + actx.from_numpy(4*np.ones((3, )))
-    qux = actx.from_numpy(4*np.ones((3, ))) + foo
+    mystate = ArrayContainerWithNumpy(
+            u=np.zeros(10),
+            v=DOFArray(actx, (actx.from_numpy(np.zeros(42)),)),
+            )
 
-    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
-                               actx.to_numpy(baz.u[0]))
+    from arraycontext import rec_map_array_container
+    rec_map_array_container(lambda x: x, mystate)
 
-    np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
-                               actx.to_numpy(qux.u[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)
+# }}}
+
+
+# {{{ 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__}'")
 
-    if _actx_allows_scalar_broadcast(actx):
-        quux = foo + actx.from_numpy(np.array(4))
-        quuz = actx.from_numpy(np.array(4)) + foo
+    import pyopencl.array as cl_array
+    ary = cl_array.to_device(actx.queue, np.zeros((32, 7)))
 
-        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
-                                   actx.to_numpy(quux.u[0]))
+    # {{{ check tags are set
 
-        np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
-                                   actx.to_numpy(quuz.u[0]))
+    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 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)
+
+
+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"
 
@@ -22,11 +26,18 @@ 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
 
@@ -36,13 +47,193 @@ def test_pt_actx_key_stringification_uniqueness():
     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:
-        from pytest import main
-        main([__file__])
+        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)