Repository: google/jax
Pull request #21410 (2024-05-24 00:21:08)

Add the `threefry_partitionable` config to `JaxprEqnContext` to allow setting it inside jit.
Add the `threefry_partitionable` config to `JaxprEqnContext` to allow setting it inside jit.

Before this information was lost in the roundtrip via `mlir.lower_fun` -> `jaxpr_subcomp`. But now since it's on the jaxpr equations, the information is preserved in jaxpr_subcomp as we enter into each eqn's ctx.

Fixes: https://github.com/google/jax/issues/21061



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 3bd9b14a9..806ffeefc 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -173,7 +173,6 @@ class Jaxpr:
     return jaxpr
 
 
-
 def join_effects(*effects: Effects) -> Effects:
   return set().union(*effects) if effects else no_effects
 
@@ -262,9 +261,13 @@ def jaxpr_as_fun(closed_jaxpr: ClosedJaxpr, *args):
 
 class JaxprEqnContext:
 
-  def __init__(self, compute_type: str | None):
+  def __init__(self, compute_type: str | None, threefry_partitionable: bool):
     self.compute_type = compute_type
-    self._managers = [(compute_on.extend_compute_type, self.compute_type)]
+    self.threefry_partitionable = threefry_partitionable
+    self._managers = [
+        (compute_on.extend_compute_type, self.compute_type),
+        (config.threefry_partitionable.__call__, self.threefry_partitionable),
+    ]
 
   @property
   @contextmanager
@@ -275,7 +278,8 @@ class JaxprEqnContext:
       yield
 
   def __repr__(self):
-    return f'JaxprEqnContext(compute_type={self.compute_type})'
+    return (f"JaxprEqnContext(compute_type={self.compute_type},"
+            f"threefry_partitionable={self.threefry_partitionable})")
 
 
 class JaxprEqn(NamedTuple):
@@ -316,7 +320,8 @@ class JaxprEqn(NamedTuple):
 def new_jaxpr_eqn(invars, outvars, primitive, params, effects, source_info=None,
                   ctx=None):
   source_info = source_info or source_info_util.new_source_info()
-  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type())
+  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type(),
+                               config.threefry_partitionable.value)
   if config.enable_checks.value:
     assert all(isinstance(x, (Var, Literal)) for x in  invars)
     assert all(isinstance(v,  Var)           for v in outvars)
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 600f10e9f..36d3c884b 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1547,7 +1547,8 @@ def jaxpr_subcomp(ctx: ModuleContext, jaxpr: core.Jaxpr,
     source_info = eqn.source_info.replace(
         name_stack=name_stack + eqn.source_info.name_stack)
     loc = _source_info_to_location(ctx, eqn.primitive, eqn.params, source_info)
-    with source_info_util.user_context(eqn.source_info.traceback), loc:
+    with (source_info_util.user_context(eqn.source_info.traceback), loc,
+          eqn.ctx.manager):
       override_rule = get_override_lowering_rule(eqn.primitive)
       platform_rules: dict[str, LoweringRule] = {}
       default_rule: LoweringRule | None = None
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 64c4c0009..704b55578 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -897,7 +897,8 @@ def new_eqn_recipe(in_tracers: Sequence[JaxprTracer],
     assert ("donated_invars" in params and
             len(params["donated_invars"]) == len(params["call_jaxpr"].invars))
   out_avals = [core.raise_to_shaped(t.aval) for t in out_tracers]
-  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type())
+  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type(),
+                               config.threefry_partitionable.value)
   return JaxprEqnRecipe(object(), tuple(in_tracers), map(ref, out_tracers),
                         out_avals, primitive, params, effects, source_info,
                         ctx)
@@ -1259,7 +1260,8 @@ def _partial_eval_jaxpr_custom_cached(
         offload_eqn = core.JaxprEqn(
             outvars_copy, resvars, device_put_p,
             dict(device=TransferToMemoryKind(policy.dst), src=None),
-            set(), source_info_util.new_source_info(), JaxprEqnContext(None))
+            set(), source_info_util.new_source_info(),
+            JaxprEqnContext(None, False))
         known_eqns.append(offload_eqn)
         # resvars are known and available in the backward jaxpr.
         map(partial(write, False, True), resvars)
@@ -1267,7 +1269,8 @@ def _partial_eval_jaxpr_custom_cached(
         reload_eqn = core.JaxprEqn(
             resvars, eqn.outvars, device_put_p,  # type: ignore
             dict(device=TransferToMemoryKind(policy.src), src=None),
-            set(), source_info_util.new_source_info(), JaxprEqnContext(None))
+            set(), source_info_util.new_source_info(),
+            JaxprEqnContext(None, False))
         staged_eqns.append(reload_eqn)
         # outvars are known and available in the backward jaxpr.
         map(partial(write, False, True), eqn.outvars)


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 23bcbfded..7bb79f212 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1253,18 +1253,18 @@ class RngShardingTest(jtu.JaxTestCase):
     x = array.make_array_from_callback(global_x.shape, s, lambda i: global_x[i])
 
     # check computation is fully partitioned and without any communication
-    jax.config.update('jax_threefry_partitionable', True)
-    unopt_txt = f.lower(x).as_text(dialect='hlo')
-    opt_txt = f.lower(x).compile().as_text()
-    self.assertIn(   f'[{n}]', unopt_txt)
-    self.assertNotIn(f'[{n}]', opt_txt)
-    self.assertNotIn('all-reduce', opt_txt)
-    self.assertNotIn('collective-permute', opt_txt)
-
-    # check against single-device reference
-    y = f(x)
-    y_ref1 = f(jax.device_put(x, jax.devices()[0]))
-    self.assertArraysEqual(y, y_ref1)
+    with jax.threefry_partitionable(True):
+      unopt_txt = f.lower(x).as_text(dialect='hlo')
+      opt_txt = f.lower(x).compile().as_text()
+      self.assertIn(   f'[{n}]', unopt_txt)
+      self.assertNotIn(f'[{n}]', opt_txt)
+      self.assertNotIn('all-reduce', opt_txt)
+      self.assertNotIn('collective-permute', opt_txt)
+
+      # check against single-device reference
+      y = f(x)
+      y_ref1 = f(jax.device_put(x, jax.devices()[0]))
+      self.assertArraysEqual(y, y_ref1)
 
   @parameterized.named_parameters(
       {"testcase_name": f"_{mesh_shape}_{pspec}",
@@ -1289,19 +1289,19 @@ class RngShardingTest(jtu.JaxTestCase):
     x = array.make_array_from_callback(global_x.shape, s, lambda i: global_x[i])
 
     # check computation is fully partitioned and without any communication
-    jax.config.update('jax_threefry_partitionable', True)
-    unopt_txt = f.lower(x).as_text(dialect='hlo')
-    opt_txt = f.lower(x).compile().as_text()
-    global_shape_fmt = ','.join(str(x) for x in global_shape)
-    self.assertIn(   f'[{global_shape_fmt}]', unopt_txt)
-    self.assertNotIn(f'[{global_shape_fmt}]', opt_txt)
-    self.assertNotIn('all-reduce', opt_txt)
-    self.assertNotIn('collective-permute', opt_txt)
-
-    # check against single-device reference
-    y = f(x)
-    y_ref1 = f(jax.device_put(x, jax.devices()[0]))
-    self.assertArraysEqual(y, y_ref1)
+    with jax.threefry_partitionable(True):
+      unopt_txt = f.lower(x).as_text(dialect='hlo')
+      opt_txt = f.lower(x).compile().as_text()
+      global_shape_fmt = ','.join(str(x) for x in global_shape)
+      self.assertIn(   f'[{global_shape_fmt}]', unopt_txt)
+      self.assertNotIn(f'[{global_shape_fmt}]', opt_txt)
+      self.assertNotIn('all-reduce', opt_txt)
+      self.assertNotIn('collective-permute', opt_txt)
+
+      # check against single-device reference
+      y = f(x)
+      y_ref1 = f(jax.device_put(x, jax.devices()[0]))
+      self.assertArraysEqual(y, y_ref1)
 
 
 if __name__ == '__main__':
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 208e272a1..1119b09bb 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4123,6 +4123,20 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertArraysEqual(expected_y.max(), y.max())
     self.assertArraysEqual(expected_z.max(), z.max())
 
+  def test_threefry_partitionable_context_within_jit(self):
+    with jax.threefry_partitionable(False):
+      def f(x):
+        return x + jax.random.randint(jax.random.key(72), (), 0, 10)
+
+      def g(x):
+        with jax.threefry_partitionable(True):  # False by default
+          return x + jax.random.randint(jax.random.key(72), (), 0, 10)
+
+      h = jax.jit(g)
+
+      self.assertNotEqual(f(1), g(1))
+      self.assertEqual(g(1), h(1))
+
 
 class TempSharding(Sharding):