add MPIEagerJaxArrayContext

Author: matthiasdiener

grudge/array_context.py

@@ -76,10 +76,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 +429,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 +542,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)
 
 # }}}
 

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_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