Thread PRNG key through wrap_jax_jit to prevent UnexpectedTracerError

Random aten ops (dropout, bernoulli, randn, …) call
env.get_and_rotate_prng_key() which mutates RuntimeProperty.prng to a
DynamicJaxprTracer while jit-tracing. The tracer escapes the trace
and contaminates env.prng_key; subsequent reads return a tracer
instead of a usable key, and any later jax.jit invocation that
re-traces and references the leaked tracer raises
jax.errors.UnexpectedTracerError.

This fixes the leak by threading the PRNG key through the jit
boundary as the last positional input and trailing output. Random
ops rotate on a scoped RuntimeProperty (popped on trace exit), the
rotated key is returned as a jit output, and the env is refreshed
outside jit so subsequent calls see the advanced key. The fix is
transparent to callers — wrap_jax_jit's signature is unchanged.

Also adds a prng_key setter (the only mutation path was previously
manual_seed) so wrap_jax_jit can refresh the env after each jit
call.

Closes #17.

Signed-off-by: Fei Ding <caffeding@gmail.com>

Author: ifding

test/test_interop.py

@@ -179,5 +179,54 @@ def test_torch_jax_view_dtype(self):
     self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
 
 
+class JaxJitPRNGThreadingTest(unittest.TestCase):
+  """Regression tests for PRNG-key threading through `wrap_jax_jit`. Closes #17."""
+
+  def setUp(self):
+    torchax.enable_globally()
+    self.env = torchax.default_env()
+    dropout = torch.nn.Dropout(p=0.5).train()
+    self.jitted = interop.jax_jit(lambda x: dropout(x))
+    self.x = torch.ones(64, device="jax")
+
+  def test_random_op_does_not_leak_tracer(self):
+    """One jit-wrapped call with a random op must not leak a tracer into env.prng_key."""
+    self.jitted(self.x)
+    self.assertFalse(
+      isinstance(self.env.prng_key, jax.core.Tracer),
+      f"env.prng_key leaked a {type(self.env.prng_key).__name__}",
+    )
+
+  def test_consecutive_calls_advance_prng(self):
+    """Two same-input calls must differ — the threaded prng key advances."""
+    a = self.jitted(self.x)
+    b = self.jitted(self.x)
+    self.assertFalse(torch.allclose(a, b), "PRNG key not rotating between calls")
+
+  def test_setter_symmetric_with_getter(self):
+    """`env.prng_key = key` round-trips through the getter."""
+    new_key = jax.random.PRNGKey(0xDEADBEEF)
+    self.env.prng_key = new_key
+    self.assertTrue(
+      jax.numpy.array_equal(
+        jax.random.key_data(self.env.prng_key), jax.random.key_data(new_key)
+      )
+    )
+
+  def test_manual_seed_controls_output(self):
+    """Same seed → same outputs; different seed → different; successive calls advance."""
+    self.env.manual_seed(42)
+    a1, a2 = self.jitted(self.x), self.jitted(self.x)
+    self.env.manual_seed(42)
+    b1, b2 = self.jitted(self.x), self.jitted(self.x)
+    self.env.manual_seed(99)
+    c1 = self.jitted(self.x)
+
+    torch.testing.assert_close(a1, b1)
+    torch.testing.assert_close(a2, b2)
+    self.assertFalse(torch.allclose(a1, c1), "different seeds produced same output")
+    self.assertFalse(torch.allclose(a1, a2), "PRNG did not advance within seed scope")
+
+
 if __name__ == "__main__":
   unittest.main()

torchax/interop.py

@@ -384,10 +384,32 @@ def backward(ctx, *grad_out):
 
 
 def wrap_jax_jit(torch_function, jax_jit_func=jax.jit, kwargs_for_jax=None):
+  """Jit a torch function via JAX, threading the PRNG key through the jit
+  boundary so random ops (dropout, bernoulli, …) don't leak tracers into
+  `env.prng_key`. The key is appended as the last positional input/output
+  so positional `donate_argnums` stays valid. Callers passing
+  `in_shardings`/`out_shardings`/`in_specs`/`out_specs` must extend them
+  by one replicated entry for the prng leaf.
+  """
   kwargs_for_jax = kwargs_for_jax or {}
+  env = torchax.default_env()
   jax_func = jax_view(torch_function)
-  jitted = jax_jit_func(jax_func, **kwargs_for_jax)
-  return torch_view(jitted)
+
+  def jax_func_with_prng(*args_and_key, **kwargs):
+    *args, prng_key = args_and_key
+    with env.override_property(prng=prng_key):
+      out = jax_func(*args, **kwargs)
+      new_key = env.prng_key
+    return out, new_key
+
+  jitted = jax_jit_func(jax_func_with_prng, **kwargs_for_jax)
+
+  def call_with_prng(*args, **kwargs):
+    out, new_key = jitted(*args, env.prng_key, **kwargs)
+    env.prng_key = new_key
+    return out
+
+  return torch_view(call_with_prng)
 
 
 def jax_jit(torch_function, kwargs_for_jax_jit=None, fix_for_buffer_donation=False):

torchax/tensor.py

@@ -416,6 +416,12 @@ def manual_seed(self, key):
   def prng_key(self):
     return self.param.prng
 
+  @prng_key.setter
+  def prng_key(self, key):
+    # Mirror of the getter. Use `override_property(prng=...)` for
+    # trace-scoped keys; calling this inside `jax.jit` leaks a tracer.
+    self.param.prng = key
+
   def _should_use_torchax_tensor(self, device):
     if device is None:
       device = torch.get_default_device()