Opens up API hole enabling jax_privacy to disable ability to read jax abstract mesh in drjax.

PiperOrigin-RevId: 766232793

Author: jkr26

drjax/__init__.py

@@ -32,6 +32,9 @@
 @_functools.wraps(_api.drjax_program)
 def program(**kwargs):
   """A decorator enabling calling the DrJAX API."""
+  # TODO(b/421499635): Remove this API hole once we diagnose and fix clipping
+  # ooming.
+  use_abstract_mesh = kwargs.pop('use_abstract_mesh', True)
   try:
     [(placement, value)] = kwargs.items()  # pylint: disable=unbalanced-dict-unpacking
   except ValueError as e:
@@ -50,10 +53,13 @@ def program(**kwargs):
       )
       # pylint: enable=f-string-without-interpolation
     return _api.drjax_program(
-        placements=value, self_module=_sys.modules[__name__]
+        placements=value,
+        self_module=_sys.modules[__name__],
+        use_abstract_mesh=use_abstract_mesh,
     )
   else:
     return _api.drjax_program(
         placements=kwargs,
         self_module=_sys.modules[__name__],
+        use_abstract_mesh=use_abstract_mesh,
     )

drjax/_src/api.py

@@ -233,6 +233,7 @@ def drjax_program(
     *,
     placements: Mapping[str, int],
     self_module,
+    use_abstract_mesh: bool = True,
 ):
   """Patches symbols into current module and call `jax.jit` on the result.
 
@@ -261,6 +262,9 @@ def drjax_program(
       collectives referencing this name results in undefined behavior).
     self_module: The Python module to patch the API when performing DrJAX
       tracing.
+    use_abstract_mesh: Whether to optionally search for jax's abstract mesh when
+      adding drjax sharding constraints (e.g. making use of drjax compatible
+      with jax.sharding.use_mesh).
 
   Returns:
     A decorated function enabling the calling of the DrJAX API. Interoperable
@@ -275,6 +279,7 @@ def drjax_program(
 
   placed_computations = impls.PlacedComputations(
       placements_to_n_elements=placements,
+      use_abstract_mesh=use_abstract_mesh,
   )
   prim_computations, primdefs = primitives.register_primitives(
       placements=placements

drjax/_src/impls.py

@@ -47,9 +47,13 @@ def call_jaxpr(fn, arg):
 
 # TODO(b/366437841): Remove use of pxla.thread_resources.env.physical_mesh,
 # which is a JAX internal API.
-def _global_mesh() -> jax.sharding.Mesh | jax.sharding.AbstractMesh | None:
+def _global_mesh(
+    use_abstract: bool,
+) -> jax.sharding.Mesh | jax.sharding.AbstractMesh | None:
   """Returns the JAX global mesh if installed, or `None` otherwise."""
-  jax_global_mesh = jax.sharding.get_abstract_mesh()
+  jax_global_mesh = None
+  if use_abstract:
+    jax_global_mesh = jax.sharding.get_abstract_mesh()
   if jax_global_mesh is None or jax_global_mesh.empty:
     jax_global_mesh = pxla.thread_resources.env.physical_mesh
   return None if jax_global_mesh.empty else jax_global_mesh
@@ -93,8 +97,10 @@ class PlacedComputations:
   def __init__(
       self,
       placements_to_n_elements: Mapping[str, int],
+      use_abstract_mesh: bool = True,
   ):
     self._placements_to_n_elements = placements_to_n_elements
+    self._use_abstract_mesh = use_abstract_mesh
 
   def broadcast_to_placement(
       self,
@@ -130,7 +136,7 @@ def broadcast_to_placement(
       A logically tiled array along the zeroth axis, as described above.
     """
     if mesh is None:
-      mesh = _global_mesh()
+      mesh = _global_mesh(self._use_abstract_mesh)
 
     arg = jnp.array(arg)
     n_elements = self._placements_to_n_elements[placement]
@@ -240,7 +246,7 @@ def map_to_placement(
       requirements specified above.
     """
     if mesh is None:
-      mesh = _global_mesh()
+      mesh = _global_mesh(self._use_abstract_mesh)
 
     def _constrain_at_placement_with_slices_like(x, y):
       pspec = P(placement, *([P.UNCONSTRAINED] * (len(x.shape) - 1)))