-
Couldn't load subscription status.
- Fork 128
BitGenerator support #499
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
BitGenerator support #499
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks like a useful addition! Thanks for working on it. I'm definitely not an expert here, but I left a few comment about things that stood out to me. Let me know what you think.
Also, are there any differences between numpy v1 and v2 that we need to consider?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should be removed
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
sure, will do when I’m done. I like working on multiple machines, and I don’t like re-doing settings for individual projects
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I still don't love the drop impl, but with the way to manually release it with a Python token, it may be acceptable. Maybe @davidhewitt has an idea and/or comments about the appoach. Otherwise I only have a few minor remarks.
Thanks for the comments, I’ll address them!
The main issue is that I think I’m triggering UB somehow and I don’t know how: when running all tests, often some unrelated test run after this one crashes ...
Also, are there any differences between numpy v1 and v2 that we need to consider?
I didn’t forget about this either, will look!
/edit: the C API for random is there since 1.19: https://numpy.org/doc/1.26/reference/random/c-api.html
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It feels like there should be a convenient way to get this. I'm thinking about something like
impl PyBitGenerator { fn new(py: Python<'_>) -> PyResult<Bound<..>>; }
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
there are many implementations, we’d have to cover all of them.
I’d rather leave this minimal until this PR is mostly done.
src/random.rs
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This may not be true under free-threaded Python. Is the lock known to be threadsafe and acquire simply fails if the lock is already acquired? If not we may need to guard the whole module under cfg(not(Py_GIL_DISABLED))
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it doesn’t fail, it hangs, but that’s configurable with a timeout or by making it non-blocking: https://docs.python.org/3/library/threading.html#threading.Lock.acquire
and it’s a threading.Lock!
OK, with the release attr and changing the parallel test to use the explicit release as well, the UB now sometimes manifests as a lock poisoning error. progress?
I may have found a problem:
This fails as intended:
Python::with_gil(|py| { let obj = get_bit_generator(py)?; let a = obj.lock()?; let b = obj.lock()?; Ok::<_, PyErr>(()) }) .unwrap();
returning
called `Result::unwrap()` on an `Err` value: PyErr { type: <class 'RuntimeError'>, value: RuntimeError('BitGenerator is already locked'), traceback: None }
But this does not fail:
Python::with_gil(|py| { let a = get_bit_generator(py)?.lock()?; let b = get_bit_generator(py)?.lock()?; Ok::<_, PyErr>(()) }) .unwrap();
and crucially it gives the same pointers:
[src/random.rs:113:18] ptr = 0x00007b9f6be44cc0
[src/random.rs:113:18] *ptr = bitgen_t {
state: 0x00007b9f6be44d08,
next_uint64: 0x00007b9f6837d320,
next_uint32: 0x00007b9f6837d370,
next_double: 0x00007b9f6837d3f0,
next_raw: 0x00007b9f6837d320,
}
[src/random.rs:113:18] ptr = 0x00007b9f6be44cc0
[src/random.rs:113:18] *ptr = bitgen_t {
state: 0x00007b9f6be44d08,
next_uint64: 0x00007b9f6837d320,
next_uint32: 0x00007b9f6837d370,
next_double: 0x00007b9f6837d3f0,
next_raw: 0x00007b9f6837d320,
}
So when using multiple threads, for example multiple tests running in parallel, we have a data race on the state. I think we need a lock across all instances to make this work.
Oh wow, so while default_rng(...).bit_generator.state is always different (somehow), default_rng(...).bit_generator.ctypes.state_address isn’t necessarily (somehow).
note the different seed sequence passed to the function, not even then is the state address different wtf:
>>> np.random.default_rng([1, 4]).bit_generator.ctypes.state_address 4355856392 >>> np.random.default_rng([2, 4]).bit_generator.ctypes.state_address 4355856392
I have no clue what to make of that. I just assumed different random state on the Python side means a different underlying struct, because how can that not be the case?
But anyway, you made me realize that the whole approach is flawed because the same BitGenerator can be passed from the Python side multiple times. So a generator passed from Python doesn’t have guaranteed independent state from another. Therefore if we want to use it, we’d have to use its threading lock as intended instead of abusing that lock into meaning "we can now do whatever we want with it"
So I think the way to go is instead of locking to use spawn to get independent child generators (which are always different, so we could use them to our hearts’ content, but also one should probably use one per thread anyway):
>>> [bg.ctypes.state_address for bg in np.random.default_rng().bit_generator.spawn(2)] [4355860968, 4355862376]
Maybe we should skip the guard part and just lock and unlock within the RngCore implementation itself. Can you give an example for why you'd want this, and why the api has this form? Why would someone want to use this rather than the RngCore impl that rand ships with? Maybe we can come up with a better design.
When implementing Python-facing APIs, having a rng: np.random.Generator parameter is common. I want to write code that actually respects that parameter and uses it instead of ignoring it or calling it once to seed the actual generator.
juntyr
commented
Jun 24, 2025
Would it also possible to go the other way, i.e. provide a rand rng from Rust to Python as a numpy BitGenerator?
Yes, that's part of numpy's API as well!
Uh oh!
There was an error while loading. Please reload this page.
See
Fixes #498
The idea is to have a safe wrapper around the
npy_bitgenstruct that implementsrand::RngCore. That way pyo3 functions could be passed anp.random.Generator, get that wrapper from it, and pass it to Rust APIs, which could then call its methods repeatedly.The way it’s implemented, the workflow would look like this:
downcastanp.random.BitGeneratorinstance into anumpy::random::PyBitGenerator..lock()on it to get anumpy::random::PyBitGeneratorGuard.TODO:
Safety
If somebody releases the threading lock of the
BitGeneratorwhile we’re using it, this isn’t safe 🤔API design options
I could make this more complex by adding a new trait that is implemented by both
PyBitGeneratorandPyBitGeneratorGuard, allowing to choose if someone wants toPyBitGenerator’srandom_*methods directly on that object while holding the GIL and without locking itnp.random.BitGeneratorand returning a GIL-free object that can be used.but for now I just implemented the use case that’s actually desired.