add MPIEagerJAXArrayContext

.test-conda-env-py3.yml

@@ -12,6 +12,7 @@ dependencies:
 - pyopencl
 - python=3
 - gmsh
+- jax
 
 # test scripts use ompi-specific arguments
 - openmpi

doc/conf.py

@@ -19,6 +19,7 @@
 intersphinx_mapping = {
         "arraycontext": ("https://documen.tician.de/arraycontext/", None),
         "loopy": ("https://documen.tician.de/loopy/", None),
+        "jax": ("https://docs.jax.dev/en/latest/", None),
         "meshmode": ("https://documen.tician.de/meshmode/", None),
         "modepy": ("https://documen.tician.de/modepy/", None),
         "mpi4py": ("https://mpi4py.readthedocs.io/en/stable", None),
@@ -32,6 +33,7 @@
 os.environ["PYOPENCL_TEST"] = "port:cpu"
 
 nitpick_ignore_regex = [
+    ["py:mod", r"jax"],  # FIXME: not sure why this does not work
     ["py:class", r"np\.ndarray"],
     ["py:data|py:class", r"arraycontext.*ContainerTc"],
 ]

grudge/array_context.py

@@ -5,6 +5,7 @@
 .. autoclass:: MPIPyOpenCLArrayContext
 .. autoclass:: MPINumpyArrayContext
 .. class:: MPIPytatoArrayContext
+.. autoclass:: MPIEagerJAXArrayContext
 .. autofunction:: get_reasonable_array_context_class
 """
 
@@ -76,10 +77,11 @@
     _HAVE_FUSION_ACTX = False
 
 
-from arraycontext import ArrayContext, NumpyArrayContext
+from arraycontext import ArrayContext, EagerJAXArrayContext, NumpyArrayContext
 from arraycontext.container import ArrayContainer
 from arraycontext.impl.pytato.compile import LazilyPyOpenCLCompilingFunctionCaller
 from arraycontext.pytest import (
+    _PytestEagerJaxArrayContextFactory,
     _PytestNumpyArrayContextFactory,
     _PytestPyOpenCLArrayContextFactoryWithClass,
     _PytestPytatoPyOpenCLArrayContextFactory,
@@ -428,6 +430,26 @@ def clone(self) -> Self:
 # }}}
 
 
+# {{{ distributed + eager jax
+
+class MPIEagerJAXArrayContext(EagerJAXArrayContext, MPIBasedArrayContext):
+    """An array context for using distributed computation with :mod:`jax`
+    eager evaluation.
+
+    .. autofunction:: __init__
+    """
+
+    def __init__(self, mpi_communicator) -> None:
+        super().__init__()
+
+        self.mpi_communicator = mpi_communicator
+
+    def clone(self) -> Self:
+        return type(self)(self.mpi_communicator)
+
+# }}}
+
+
 # {{{ distributed + pytato array context subclasses
 
 class MPIBasePytatoPyOpenCLArrayContext(
@@ -521,12 +543,23 @@ def __call__(self):
         return self.actx_class()
 
 
+class PytestEagerJAXArrayContextFactory(_PytestEagerJaxArrayContextFactory):
+    actx_class = EagerJAXArrayContext
+
+    def __call__(self):
+        import jax
+        jax.config.update("jax_enable_x64", True)
+        return self.actx_class()
+
+
 register_pytest_array_context_factory("grudge.pyopencl",
         PytestPyOpenCLArrayContextFactory)
 register_pytest_array_context_factory("grudge.pytato-pyopencl",
         PytestPytatoPyOpenCLArrayContextFactory)
 register_pytest_array_context_factory("grudge.numpy",
         PytestNumpyArrayContextFactory)
+register_pytest_array_context_factory("grudge.eager-jax",
+        PytestEagerJAXArrayContextFactory)
 
 # }}}
 

grudge/geometry/metrics.py

@@ -566,15 +566,21 @@ def _signed_face_ones(
 
     signed_face_ones_numpy = actx.to_numpy(signed_ones)
 
+    new_group_arrays = []
     for igrp, grp in enumerate(all_faces_conn.groups):
+        grp_field = signed_face_ones_numpy[igrp]
+        sign_mask = np.ones_like(grp_field)
+
         for batch in grp.batches:
             assert batch.to_element_face is not None
             i = actx.to_numpy(actx.thaw(batch.to_element_indices))
-            grp_field = signed_face_ones_numpy[igrp].reshape(-1)
-            grp_field[i] = \
-                (2.0 * (batch.to_element_face % 2) - 1.0) * grp_field[i]
+            sign = (2.0 * (batch.to_element_face % 2) - 1.0)
+            sign_mask[i, :] = sign
+
+        new_group_arrays.append(grp_field * sign_mask)
 
-    return actx.from_numpy(signed_face_ones_numpy)
+    from meshmode.dof_array import DOFArray
+    return actx.from_numpy(DOFArray(actx, tuple(new_group_arrays)))
 
 
 def parametrization_derivative(

test/test_dt_utils.py

@@ -27,6 +27,7 @@
 from arraycontext import pytest_generate_tests_for_array_contexts
 
 from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
     PytestNumpyArrayContextFactory,
     PytestPyOpenCLArrayContextFactory,
     PytestPytatoPyOpenCLArrayContextFactory,
@@ -36,7 +37,8 @@
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
         [PytestPyOpenCLArrayContextFactory,
          PytestPytatoPyOpenCLArrayContextFactory,
-         PytestNumpyArrayContextFactory])
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 import logging
 

test/test_euler_model.py

@@ -31,12 +31,18 @@
 )
 
 from grudge import op
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 
 
 logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 @pytest.mark.parametrize("order", [1, 2, 3])

test/test_grudge.py

@@ -42,13 +42,19 @@
 from pytools.obj_array import flat_obj_array
 
 from grudge import dof_desc, geometry, op
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 from grudge.discretization import make_discretization_collection
 
 
 logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 # {{{ mass operator trig integration

test/test_metrics.py

@@ -33,6 +33,7 @@
 from meshmode.dof_array import flat_norm
 
 from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
     PytestNumpyArrayContextFactory,
     PytestPyOpenCLArrayContextFactory,
     PytestPytatoPyOpenCLArrayContextFactory,
@@ -44,7 +45,8 @@
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
         [PytestPyOpenCLArrayContextFactory,
          PytestPytatoPyOpenCLArrayContextFactory,
-         PytestNumpyArrayContextFactory])
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 # {{{ inverse metric

test/test_modal_connections.py

@@ -23,12 +23,18 @@
 
 from arraycontext import pytest_generate_tests_for_array_contexts
 
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 from grudge.discretization import make_discretization_collection
 
 
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 import pytest
 

test/test_mpi_communication.py

@@ -35,7 +35,12 @@
 from pytools.obj_array import flat_obj_array
 
 from grudge import dof_desc, op
-from grudge.array_context import MPIPyOpenCLArrayContext, MPIPytatoArrayContext
+from grudge.array_context import (
+    MPIEagerJAXArrayContext,
+    MPINumpyArrayContext,
+    MPIPyOpenCLArrayContext,
+    MPIPytatoArrayContext,
+)
 from grudge.discretization import make_discretization_collection
 from grudge.shortcuts import compiled_lsrk45_step
 
@@ -49,7 +54,8 @@ class SimpleTag:
 
 # {{{ mpi test infrastructure
 
-DISTRIBUTED_ACTXS = [MPIPyOpenCLArrayContext, MPIPytatoArrayContext]
+DISTRIBUTED_ACTXS = [MPIPyOpenCLArrayContext, MPIPytatoArrayContext,
+                     MPIEagerJAXArrayContext, MPINumpyArrayContext]
 
 
 def run_test_with_mpi(num_ranks, f, *args):
@@ -87,6 +93,10 @@ def run_test_with_mpi_inner():
         actx = actx_class(comm, queue, mpi_base_tag=15000)
     elif actx_class is MPIPyOpenCLArrayContext:
         actx = actx_class(comm, queue)
+    elif actx_class is MPIEagerJAXArrayContext:
+        actx = actx_class(comm)
+    elif actx_class is MPINumpyArrayContext:
+        actx = actx_class(comm)
     else:
         raise ValueError("unknown actx_class")
 

test/test_op.py

@@ -36,7 +36,11 @@
 from pytools.obj_array import make_obj_array
 
 from grudge import geometry, op
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 from grudge.discretization import make_discretization_collection
 from grudge.dof_desc import (
     DISCR_TAG_BASE,
@@ -52,7 +56,9 @@
 
 logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 # {{{ gradient

test/test_reductions.py

@@ -39,13 +39,19 @@
 from pytools.obj_array import make_obj_array
 
 from grudge import op
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 from grudge.discretization import make_discretization_collection
 
 
 logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 @pytest.mark.parametrize(("mesh_size", "with_initial"), [

test/test_tools.py

@@ -26,25 +26,11 @@
 
 import numpy as np
 import numpy.linalg as la  # noqa
-
-from arraycontext import pytest_generate_tests_for_array_contexts
-
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
-
-
-pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
-
-import logging
-
 import pytest
 
 from pytools.obj_array import make_obj_array
 
 
-logger = logging.getLogger(__name__)
-
-
 # {{{ map_subarrays and rec_map_subarrays
 
 @dataclass(frozen=True, eq=True)

test/test_trace_pair.py

@@ -29,14 +29,20 @@
 from arraycontext import pytest_generate_tests_for_array_contexts
 from meshmode.dof_array import DOFArray
 
-from grudge.array_context import PytestPyOpenCLArrayContextFactory
+from grudge.array_context import (
+    PytestEagerJAXArrayContextFactory,
+    PytestNumpyArrayContextFactory,
+    PytestPyOpenCLArrayContextFactory,
+)
 from grudge.discretization import make_discretization_collection
 from grudge.trace_pair import TracePair
 
 
 logger = logging.getLogger(__name__)
 pytest_generate_tests = pytest_generate_tests_for_array_contexts(
-        [PytestPyOpenCLArrayContextFactory])
+        [PytestPyOpenCLArrayContextFactory,
+         PytestNumpyArrayContextFactory,
+         PytestEagerJAXArrayContextFactory])
 
 
 def test_trace_pair(actx_factory):