Repository: google/jax

Issue:
`make_jaxpr` does not work with any parallel operators
Currently `make_jaxpr` does not work when there are any parallel operations, such as `psum/pmean` inside the computation. This also is reflected in the fact that the in general the returned `jaxpr` returned by `make_jaxpr` is not aware of any parallel computation in general.

This breaks any code that needs to analyze or internally transform functions based on their `jaxpr`. 
`make_jaxpr` needs some way of adding the parallel environment to it.


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -2000,6 +2000,7 @@ def transposed_fun(out_cotangent):
 
 def make_jaxpr(fun: Callable,
                static_argnums: Union[int, Iterable[int]] = (),
+               axis_env: Optional[Sequence[Tuple[AxisName, int]]] = None,
                return_shape: bool = False,
                ) -> Callable[..., core.ClosedJaxpr]:
   """Creates a function that produces its jaxpr given example args.
@@ -2009,6 +2010,12 @@ def make_jaxpr(fun: Callable,
       arguments and return value should be arrays, scalars, or standard Python
       containers (tuple/list/dict) thereof.
     static_argnums: See the :py:func:`jax.jit` docstring.
+    axis_env: Optional, a sequence of pairs where the first element is an axis
+      name and the second element is a positive integer representing the size of
+      the mapped axis with that name. This parameter is useful when lowering
+      functions that involve parallel communication collectives, and it
+      specifies the axis name/size environment that would be set up by
+      applications of :py:func:`jax.pmap`.
     return_shape: Optional boolean, defaults to ``False``. If ``True``, the
       wrapped function returns a pair where the first element is the ``jaxpr``
       and the second element is a pytree with the same structure as
@@ -2069,8 +2076,14 @@ def jaxpr_maker(*args, **kwargs):
     jaxtree_fun, out_tree = flatten_fun(wrapped, in_tree)
     in_avals = [raise_to_shaped(core.get_aval(x)) for x in jax_args]
     if config.omnistaging_enabled:
-      jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(jaxtree_fun, in_avals)
+      with ExitStack() as stack:
+        for axis_name, size in axis_env or []:
+          stack.enter_context(core.extend_axis_env(axis_name, size, None))
+        jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(jaxtree_fun, in_avals)
     else:
+      if axis_env:
+        raise NotImplementedError(
+            "axis_env argument to make_jaxpr only supported with omnistaging.")
       in_pvals = [pe.PartialVal.unknown(a) for a in in_avals]
       jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
           jaxtree_fun, in_pvals, instantiate=True, stage_out=True)  # type: ignore


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2775,6 +2775,15 @@ def test_make_jaxpr_return_shape(self):
                 api.ShapeDtypeStruct(shape=(2,), dtype=jnp.float32))
     self.assertEqual(shape_tree, expected)
 
+  def test_make_jaxpr_axis_env(self):
+    if not config.omnistaging_enabled:
+      raise unittest.SkipTest("test only works with omnistaging")
+
+    def f(x):
+      return x - lax.psum(x, 'i')
+    jaxpr = api.make_jaxpr(f, axis_env=[('i', 4)])(2)
+    self.assertIn('psum', str(jaxpr))
+
 
 class LazyTest(jtu.JaxTestCase):