add MPIEagerJaxArrayContext

Author: matthiasdiener

grudge/array_context.py

@@ -427,6 +427,27 @@ def clone(self) -> Self:
 
 # }}}
 
+# {{{ distributed +  eager jax
+
+from arraycontext.impl.jax import EagerJAXArrayContext
+
+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