Fix captured random state in compile

python/src/random.cpp

@@ -9,6 +9,7 @@
 
 #include "mlx/ops.h"
 #include "mlx/random.h"
+#include "python/src/random.h"
 #include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
@@ -63,15 +64,57 @@ PyKeySequence& default_key() {
   return ks;
 }
 
+// A process-global sentinel for `mx.random.state`. Since it is the same object
+// on every thread, capturing it (e.g. with `mx.compile`) is thread-independent;
+// the pytree traversal in trees.cpp resolves it to the calling thread's key.
+class RandomState {};
+
+nb::object random_state_sentinel() {
+  static nb::object sentinel;
+  if (sentinel.ptr() == nullptr) {
+    sentinel = nb::cast(RandomState{});
+    sentinel.inc_ref();
+  }
+  return sentinel;
+}
+
+bool is_random_state(nb::handle obj) {
+  return obj.ptr() == random_state_sentinel().ptr();
+}
+
+mx::array random_state_key() {
+  return nb::cast<mx::array>(default_key().state()[0]);
+}
+
+void set_random_state_key(const mx::array& key) {
+  default_key().state()[0] = nb::cast(key);
+}
+
 void init_random(nb::module_& parent_module) {
   auto m = parent_module.def_submodule(
       "random",
       "mlx.core.random: functionality related to random number generation");
 
+  nb::class_<RandomState>(m, "_RandomState")
+      .def("__len__", [](const RandomState&) { return 1; })
+      .def(
+          "__getitem__",
+          [](const RandomState&, int i) -> nb::object {
+            if (i != 0 && i != -1) {
+              throw nb::index_error("random state index out of range");
+            }
+            return default_key().state()[0];
+          },
+          "index"_a)
+      .def("__iter__", [](const RandomState&) {
+        return nb::iter(default_key().state());
+      });
+
   m.def("__getattr__", [&](nb::handle key) -> nb::object {
     // Create random.state lazily to avoid initializing device during import.
     if (nb::isinstance<nb::str>(key) && nb::cast<std::string>(key) == "state") {
-      return default_key().state();
+      default_key().state();
+      return random_state_sentinel();
     }
     return nb::steal(PyErr_Format(
         PyExc_AttributeError,

python/src/random.h

@@ -0,0 +1,20 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <nanobind/nanobind.h>
+
+#include "mlx/array.h"
+
+namespace mx = mlx::core;
+namespace nb = nanobind;
+
+// The process-global `mx.random.state` sentinel.
+nb::object random_state_sentinel();
+
+// True if `obj` is the RNG state sentinel.
+bool is_random_state(nb::handle obj);
+
+// Read/write the calling thread's current PRNG key.
+mx::array random_state_key();
+void set_random_state_key(const mx::array& key);

python/src/trees.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "python/src/trees.h"
+#include "python/src/random.h"
 
 template <typename T, typename U, typename V>
 void validate_subtrees(const std::vector<nb::object>& subtrees) {
@@ -153,7 +154,10 @@ void tree_visit(
 void tree_visit(nb::handle tree, std::function<void(nb::handle)> visitor) {
   std::function<void(nb::handle)> recurse;
   recurse = [&](nb::handle subtree) {
-    if (nb::isinstance<nb::list>(subtree) ||
+    if (is_random_state(subtree)) {
+      visitor(nb::cast(random_state_key()));
+    } else if (
+        nb::isinstance<nb::list>(subtree) ||
         nb::isinstance<nb::tuple>(subtree)) {
       for (auto item : subtree) {
         recurse(item);
@@ -175,7 +179,12 @@ void tree_visit_update(
     std::function<nb::object(nb::handle)> visitor) {
   std::function<nb::object(nb::handle)> recurse;
   recurse = [&](nb::handle subtree) {
-    if (nb::isinstance<nb::list>(subtree)) {
+    if (is_random_state(subtree)) {
+      // Read/write the calling thread's key; keep the sentinel in the tree.
+      set_random_state_key(
+          nb::cast<mx::array>(visitor(nb::cast(random_state_key()))));
+      return nb::cast<nb::object>(subtree);
+    } else if (nb::isinstance<nb::list>(subtree)) {
       auto l = nb::cast<nb::list>(subtree);
       for (int i = 0; i < l.size(); ++i) {
         l[i] = recurse(l[i]);

python/tests/test_compile.py

@@ -439,6 +439,74 @@ def fun():
 
         self.assertFalse(mx.allclose(fun(), fun(), 1e-2, 1e-2))
 
+    def test_compile_rng_across_threads(self):
+        # A function compiled with inputs/outputs=mx.random.state on one thread
+        # must still use (and advance/seed) the calling thread's RNG state when
+        # invoked from another thread, whether captured directly or nested.
+
+        # The state sentinel is a single global object shared across threads.
+        state_from_thread = {}
+
+        def grab():
+            state_from_thread["s"] = mx.random.state
+
+        t = threading.Thread(target=grab)
+        t.start()
+        t.join()
+        self.assertIs(mx.random.state, state_from_thread["s"])
+
+        direct = partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)(
+            lambda: mx.random.uniform(shape=(10, 10))
+        )
+
+        nested_state = [{"unused": mx.array(0.0)}, mx.random.state]
+        nested = partial(mx.compile, inputs=nested_state, outputs=nested_state)(
+            lambda: mx.random.uniform(shape=(10, 10))
+        )
+
+        for fun in (direct, nested):
+            results = {}
+
+            def worker():
+                with mx.stream(mx.cpu):
+                    a = fun()
+                    b = fun()
+                    results["advances"] = not bool(mx.allclose(a, b, 1e-2, 1e-2).item())
+                    mx.random.seed(42)
+                    c = fun()
+                    mx.random.seed(42)
+                    d = fun()
+                    results["seed_reproducible"] = bool(mx.allclose(c, d).item())
+                    mx.random.seed(1234)
+                    e = fun()
+                    results["seed_changes"] = not bool(
+                        mx.allclose(c, e, 1e-2, 1e-2).item()
+                    )
+
+            t = threading.Thread(target=worker)
+            t.start()
+            t.join()
+
+            self.assertTrue(results["advances"])
+            self.assertTrue(results["seed_reproducible"])
+            self.assertTrue(results["seed_changes"])
+
+    def test_compile_state_capture_with_rng_updates_in_place(self):
+        # Capturing mx.random.state alongside other state via outputs= must not
+        # break in-place updates of the other captured containers.
+        counter = {"v": mx.array(0.0)}
+        state = [counter, mx.random.state]
+
+        @partial(mx.compile, inputs=state, outputs=state)
+        def step():
+            counter["v"] = counter["v"] + 1.0
+            return mx.random.uniform(shape=(2,))
+
+        for _ in range(3):
+            step()
+        mx.eval(counter["v"])
+        self.assertEqual(counter["v"].item(), 3.0)
+
     def test_compile_kwargs(self):
         @mx.compile
         def fun(x, y, z):