Skip to content

Commit ce30733

Browse files
angeloskathangeloskath
authored
Fix captured random state in compile (#3828)
1 parent b7c3dd6 commit ce30733

4 files changed

Lines changed: 144 additions & 10 deletions

File tree

‎python/src/random.cpp‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
#include "mlx/ops.h"
1111
#include "mlx/random.h"
12+
#include "python/src/random.h"
1213
#include "python/src/small_vector.h"
1314
#include "python/src/utils.h"
1415

@@ -63,15 +64,54 @@ PyKeySequence& default_key() {
6364
return ks;
6465
}
6566

67+
// A process-global sentinel for `mx.random.state`. Since it is the same object
68+
// on every thread, capturing it (e.g. with `mx.compile`) is thread-independent;
69+
// the pytree traversal in trees.cpp resolves it to the calling thread's key.
70+
class RandomState {};
71+
72+
nb::object random_state_sentinel() {
73+
static nb::object sentinel = []() {
74+
auto sentinel = nb::cast(RandomState{});
75+
sentinel.inc_ref();
76+
return sentinel;
77+
}();
78+
79+
return sentinel;
80+
}
81+
82+
mx::array random_state_key() {
83+
return nb::cast<mx::array>(default_key().state()[0]);
84+
}
85+
86+
void set_random_state_key(const mx::array& key) {
87+
default_key().state()[0] = nb::cast(key);
88+
}
89+
6690
void init_random(nb::module_& parent_module) {
6791
auto m = parent_module.def_submodule(
6892
"random",
6993
"mlx.core.random: functionality related to random number generation");
7094

95+
nb::class_<RandomState>(m, "_RandomState")
96+
.def("__len__", [](const RandomState&) { return 1; })
97+
.def(
98+
"__getitem__",
99+
[](const RandomState&, int i) -> nb::object {
100+
if (i != 0 && i != -1) {
101+
throw nb::index_error("random state index out of range");
102+
}
103+
return default_key().state()[0];
104+
},
105+
"index"_a)
106+
.def("__iter__", [](const RandomState&) {
107+
return nb::iter(default_key().state());
108+
});
109+
71110
m.def("__getattr__", [&](nb::handle key) -> nb::object {
72111
// Create random.state lazily to avoid initializing device during import.
73112
if (nb::isinstance<nb::str>(key) && nb::cast<std::string>(key) == "state") {
74-
return default_key().state();
113+
default_key().state();
114+
return random_state_sentinel();
75115
}
76116
return nb::steal(PyErr_Format(
77117
PyExc_AttributeError,

‎python/src/random.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// Copyright © 2026 Apple Inc.
2+
3+
#pragma once
4+
5+
#include <nanobind/nanobind.h>
6+
7+
#include "mlx/array.h"
8+
9+
namespace mx = mlx::core;
10+
namespace nb = nanobind;
11+
12+
// The process-global `mx.random.state` sentinel.
13+
nb::object random_state_sentinel();
14+
15+
// Read/write the calling thread's current PRNG key.
16+
mx::array random_state_key();
17+
void set_random_state_key(const mx::array& key);

‎python/src/trees.cpp‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright © 2023-2024 Apple Inc.
22

33
#include "python/src/trees.h"
4+
#include "python/src/random.h"
45

56
template <typename T, typename U, typename V>
67
void validate_subtrees(const std::vector<nb::object>& subtrees) {
@@ -152,8 +153,12 @@ void tree_visit(
152153

153154
void tree_visit(nb::handle tree, std::function<void(nb::handle)> visitor) {
154155
std::function<void(nb::handle)> recurse;
156+
auto random_state = random_state_sentinel();
155157
recurse = [&](nb::handle subtree) {
156-
if (nb::isinstance<nb::list>(subtree) ||
158+
if (subtree.is(random_state)) {
159+
visitor(nb::cast(random_state_key()));
160+
} else if (
161+
nb::isinstance<nb::list>(subtree) ||
157162
nb::isinstance<nb::tuple>(subtree)) {
158163
for (auto item : subtree) {
159164
recurse(item);
@@ -174,8 +179,14 @@ void tree_visit_update(
174179
nb::object tree,
175180
std::function<nb::object(nb::handle)> visitor) {
176181
std::function<nb::object(nb::handle)> recurse;
182+
auto random_state = random_state_sentinel();
177183
recurse = [&](nb::handle subtree) {
178-
if (nb::isinstance<nb::list>(subtree)) {
184+
if (subtree.is(random_state)) {
185+
// Read/write the calling thread's key; keep the sentinel in the tree.
186+
set_random_state_key(
187+
nb::cast<mx::array>(visitor(nb::cast(random_state_key()))));
188+
return nb::cast<nb::object>(subtree);
189+
} else if (nb::isinstance<nb::list>(subtree)) {
179190
auto l = nb::cast<nb::list>(subtree);
180191
for (int i = 0; i < l.size(); ++i) {
181192
l[i] = recurse(l[i]);
@@ -262,14 +273,12 @@ nb::object tree_unflatten(
262273
}
263274

264275
nb::object structure_sentinel() {
265-
static nb::object sentinel;
266-
267-
if (sentinel.ptr() == nullptr) {
268-
sentinel = nb::capsule(&sentinel);
269-
// probably not needed but this should make certain that we won't ever
270-
// delete the sentinel
276+
static nb::object sentinel = []() {
277+
PyObject* raw_obj = PyObject_New(PyObject, &PyBaseObject_Type);
278+
nb::object sentinel = nb::steal(raw_obj);
271279
sentinel.inc_ref();
272-
}
280+
return sentinel;
281+
}();
273282

274283
return sentinel;
275284
}

‎python/tests/test_compile.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,74 @@ def fun():
439439

440440
self.assertFalse(mx.allclose(fun(), fun(), 1e-2, 1e-2))
441441

442+
def test_compile_rng_across_threads(self):
443+
# A function compiled with inputs/outputs=mx.random.state on one thread
444+
# must still use (and advance/seed) the calling thread's RNG state when
445+
# invoked from another thread, whether captured directly or nested.
446+
447+
# The state sentinel is a single global object shared across threads.
448+
state_from_thread = {}
449+
450+
def grab():
451+
state_from_thread["s"] = mx.random.state
452+
453+
t = threading.Thread(target=grab)
454+
t.start()
455+
t.join()
456+
self.assertIs(mx.random.state, state_from_thread["s"])
457+
458+
direct = partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)(
459+
lambda: mx.random.uniform(shape=(10, 10))
460+
)
461+
462+
nested_state = [{"unused": mx.array(0.0)}, mx.random.state]
463+
nested = partial(mx.compile, inputs=nested_state, outputs=nested_state)(
464+
lambda: mx.random.uniform(shape=(10, 10))
465+
)
466+
467+
for fun in (direct, nested):
468+
results = {}
469+
470+
def worker():
471+
with mx.stream(mx.cpu):
472+
a = fun()
473+
b = fun()
474+
results["advances"] = not bool(mx.allclose(a, b, 1e-2, 1e-2).item())
475+
mx.random.seed(42)
476+
c = fun()
477+
mx.random.seed(42)
478+
d = fun()
479+
results["seed_reproducible"] = bool(mx.allclose(c, d).item())
480+
mx.random.seed(1234)
481+
e = fun()
482+
results["seed_changes"] = not bool(
483+
mx.allclose(c, e, 1e-2, 1e-2).item()
484+
)
485+
486+
t = threading.Thread(target=worker)
487+
t.start()
488+
t.join()
489+
490+
self.assertTrue(results["advances"])
491+
self.assertTrue(results["seed_reproducible"])
492+
self.assertTrue(results["seed_changes"])
493+
494+
def test_compile_state_capture_with_rng_updates_in_place(self):
495+
# Capturing mx.random.state alongside other state via outputs= must not
496+
# break in-place updates of the other captured containers.
497+
counter = {"v": mx.array(0.0)}
498+
state = [counter, mx.random.state]
499+
500+
@partial(mx.compile, inputs=state, outputs=state)
501+
def step():
502+
counter["v"] = counter["v"] + 1.0
503+
return mx.random.uniform(shape=(2,))
504+
505+
for _ in range(3):
506+
step()
507+
mx.eval(counter["v"])
508+
self.assertEqual(counter["v"].item(), 3.0)
509+
442510
def test_compile_kwargs(self):
443511
@mx.compile
444512
def fun(x, y, z):

0 commit comments

Comments
 (0)