From 117db323076a59fb0f4535482545bd034aae8a24 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Wed, 8 Jul 2026 18:07:31 -0700 Subject: [PATCH 1/3] Make mx.random.state a global sentinel resolved per-thread in pytrees --- python/src/random.cpp | 45 +++++++++++++++++++++++- python/src/random.h | 20 +++++++++++ python/src/trees.cpp | 13 +++++-- python/tests/test_compile.py | 68 ++++++++++++++++++++++++++++++++++++ 4 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 python/src/random.h diff --git a/python/src/random.cpp b/python/src/random.cpp index ddddb56eed..ea6bde1f17 100644 --- a/python/src/random.cpp +++ b/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(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_(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(key) && nb::cast(key) == "state") { - return default_key().state(); + default_key().state(); + return random_state_sentinel(); } return nb::steal(PyErr_Format( PyExc_AttributeError, diff --git a/python/src/random.h b/python/src/random.h new file mode 100644 index 0000000000..170b9f96eb --- /dev/null +++ b/python/src/random.h @@ -0,0 +1,20 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#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); diff --git a/python/src/trees.cpp b/python/src/trees.cpp index 4b9ca9e123..5b5d4af9e5 100644 --- a/python/src/trees.cpp +++ b/python/src/trees.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include "python/src/trees.h" +#include "python/src/random.h" template void validate_subtrees(const std::vector& subtrees) { @@ -153,7 +154,10 @@ void tree_visit( void tree_visit(nb::handle tree, std::function visitor) { std::function recurse; recurse = [&](nb::handle subtree) { - if (nb::isinstance(subtree) || + if (is_random_state(subtree)) { + visitor(nb::cast(random_state_key())); + } else if ( + nb::isinstance(subtree) || nb::isinstance(subtree)) { for (auto item : subtree) { recurse(item); @@ -175,7 +179,12 @@ void tree_visit_update( std::function visitor) { std::function recurse; recurse = [&](nb::handle subtree) { - if (nb::isinstance(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(visitor(nb::cast(random_state_key())))); + return nb::cast(subtree); + } else if (nb::isinstance(subtree)) { auto l = nb::cast(subtree); for (int i = 0; i < l.size(); ++i) { l[i] = recurse(l[i]); diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 35e51600e1..7a2c6b9d0d 100644 --- a/python/tests/test_compile.py +++ b/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): From 73215796d6458aee8260c13bcf3b9873873bf1a0 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 10 Jul 2026 10:58:18 -0700 Subject: [PATCH 2/3] Fix sentinel initialization --- python/src/random.cpp | 13 +++++-------- python/src/trees.cpp | 18 +++++++++--------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/python/src/random.cpp b/python/src/random.cpp index ea6bde1f17..ceebb52e7c 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -70,16 +70,13 @@ PyKeySequence& default_key() { class RandomState {}; nb::object random_state_sentinel() { - static nb::object sentinel; - if (sentinel.ptr() == nullptr) { - sentinel = nb::cast(RandomState{}); + static nb::object sentinel = []() { + auto sentinel = nb::cast(RandomState{}); sentinel.inc_ref(); - } - return sentinel; -} + return sentinel; + }(); -bool is_random_state(nb::handle obj) { - return obj.ptr() == random_state_sentinel().ptr(); + return sentinel; } mx::array random_state_key() { diff --git a/python/src/trees.cpp b/python/src/trees.cpp index 5b5d4af9e5..cf05d4e25a 100644 --- a/python/src/trees.cpp +++ b/python/src/trees.cpp @@ -153,8 +153,9 @@ void tree_visit( void tree_visit(nb::handle tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (is_random_state(subtree)) { + if (subtree.is(random_state)) { visitor(nb::cast(random_state_key())); } else if ( nb::isinstance(subtree) || @@ -178,8 +179,9 @@ void tree_visit_update( nb::object tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (is_random_state(subtree)) { + if (subtree.is(random_state)) { // Read/write the calling thread's key; keep the sentinel in the tree. set_random_state_key( nb::cast(visitor(nb::cast(random_state_key())))); @@ -271,14 +273,12 @@ nb::object tree_unflatten( } nb::object structure_sentinel() { - static nb::object sentinel; - - if (sentinel.ptr() == nullptr) { - sentinel = nb::capsule(&sentinel); - // probably not needed but this should make certain that we won't ever - // delete the sentinel + static nb::object sentinel = []() { + PyObject* raw_obj = PyObject_New(PyObject, &PyBaseObject_Type); + nb::object sentinel = nb::steal(raw_obj); sentinel.inc_ref(); - } + return sentinel; + }(); return sentinel; } From f50d146367c18f822d565c4d61ed053cf0179d44 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Fri, 10 Jul 2026 16:20:03 -0700 Subject: [PATCH 3/3] Remove is_random_state function declaration --- python/src/random.h | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/src/random.h b/python/src/random.h index 170b9f96eb..2baf9d92f1 100644 --- a/python/src/random.h +++ b/python/src/random.h @@ -12,9 +12,6 @@ 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);