Implement prng_key as a mutable array (#9305)

Author: tengyifei

torchax/pyproject.toml

@@ -40,11 +40,11 @@ classifiers = [
 path = "torchax/__init__.py"
 
 [project.optional-dependencies]
-cpu = ["jax[cpu]>=0.4.30", "jax[cpu]"]
+cpu = ["jax[cpu]>=0.6.2", "jax[cpu]"]
 # Add libtpu index `-f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html`
-tpu = ["jax[cpu]>=0.4.30", "jax[tpu]"]
-cuda = ["jax[cpu]>=0.4.30", "jax[cuda12]"]
-odml = ["jax[cpu]>=0.4.30", "jax[cpu]"]
+tpu = ["jax[cpu]>=0.6.2", "jax[tpu]"]
+cuda = ["jax[cpu]>=0.6.2", "jax[cuda12]"]
+odml = ["jax[cpu]>=0.6.2", "jax[cpu]"]
 
 [tool.hatch.build.targets.wheel]
 packages = ["torchax"]

torchax/test/test_context.py

@@ -5,7 +5,7 @@
 from torchax import tensor
 import torchax.interop
 
-xla_env = tensor.Environment()
+xla_env = torchax.default_env()
 
 
 class TestContext(unittest.TestCase):
@@ -73,15 +73,16 @@ def random_op():
       random_jit = torchax.interop.jax_jit(random_op)
       self.assertIsInstance(random_jit(), tensor.Tensor)
 
-      # Result always expected to be the same for a jitted function because seeds
-      # are baked in
-      torch.testing.assert_close(random_jit(), random_jit(), atol=0, rtol=0)
+      # If we run the JIT twice, the random values should be different.
+      with self.assertRaises(AssertionError):
+        torch.testing.assert_close(random_jit(), random_jit(), atol=0, rtol=0)
 
   def test_generator_seed(self):
     with xla_env:
       x = torch.randn(2, 3, generator=torch.Generator().manual_seed(0))
       y = torch.randn(2, 3, generator=torch.Generator().manual_seed(0))
 
+      # Values will be the same given the same seed.
       torch.testing.assert_close(x, y)
 
   def test_buffer(self):

torchax/torchax/tensor.py

@@ -1,4 +1,3 @@
-import random
 import logging
 import sys
 import contextlib
@@ -13,10 +12,11 @@
 import torch.utils._mode_utils as mode_utils
 import torch.utils._python_dispatch as torch_dispatch
 import torch.utils._pytree as torch_pytree
-from torchax.view import View, NarrowInfo
+from torchax.view import View
 from torchax import config
 from torchax.ops import mappings, ops_registry
 from torchax import amp
+from jax.experimental import mutable_array
 
 logger = logging.getLogger(__name__)
 
@@ -323,11 +323,16 @@ def __init__(self, configuration=None):
     self._manually_entered = False
     self.enabled = False
     self._jax_devices = set(["jax", "jax_cpu", "xla"])
-    self.prng_key = jax.random.key(torch.initial_seed() % (1 << 63))
+    self._prng_key = mutable_array(
+        jax.random.key(torch.initial_seed() % (1 << 63)))
     self.autocast_dtype = None
 
   def manual_seed(self, key):
-    self.prng_key = jax.random.key(key)
+    self._prng_key = mutable_array(jax.random.key(key))
+
+  @property
+  def prng_key(self):
+    return self._prng_key[...]
 
   def get_as_jax_device(self, device: Any):
     if device is None:
@@ -431,8 +436,10 @@ def get_and_rotate_prng_key(self,
                               generator: Optional[torch.Generator] = None):
     if generator is not None:
       with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
-        self.prng_key = jax.random.key(generator.initial_seed() % (2**63))
-    self.prng_key, next_key = jax.random.split(self.prng_key)
+        self._prng_key[...] = jax.random.key(generator.initial_seed() % (2**63))
+    old_key = self._prng_key[...]
+    new_prng_key, next_key = jax.random.split(old_key)
+    self._prng_key[...] = new_prng_key
     return next_key
 
   def _handle_tensor_constructor(self, func, args, kwargs):