Skip to content

feat: public RandomState API + Updatable impl#321

Merged
minghuaw merged 2 commits intooxiglade:mainfrom
nicholasjpaterno:feat/public-random-state
Feb 26, 2026
Merged

feat: public RandomState API + Updatable impl#321
minghuaw merged 2 commits intooxiglade:mainfrom
nicholasjpaterno:feat/public-random-state

Conversation

@nicholasjpaterno
Copy link
Contributor

@nicholasjpaterno nicholasjpaterno commented Feb 20, 2026

Summary

Makes RandomState public with a complete API for compiled random number generation.

  • Makes RandomState constructors public: new(), with_seed(), from_key()
  • Adds next_key() for explicit state advancement
  • Adds as_array() / as_array_mut() for direct state access
  • Implements Updatable trait for compile_with_state compatibility
  • Implements Default trait
  • Adds comprehensive tests for the new public API

This enables the Rust equivalent of Python's @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) pattern, which is essential for compiled text generation with sampling.

- Make RandomState methods public (new, with_seed, next_key, seed)
- Implement Updatable trait for compile_with_state compatibility
- Add Default impl for RandomState
- Add helper methods (from_key, as_array, as_array_mut)
- Change doc example from `ignore` to `no_run` for compile-time checking
- Fix `# Safety` to `# Note` on non-unsafe as_array_mut method
- Document Default impl panic conditions
- Add tests for new(), with_seed(), next_key(), from_key(), Updatable, Default
@nicholasjpaterno
Copy link
Contributor Author

@dcvz Do you have any thoughts on this PR and #320? Its currently a blocker for my PMetal release.

@dcvz
Copy link
Collaborator

dcvz commented Feb 24, 2026

@dcvz Do you have any thoughts on this PR and #320? Its currently a blocker for my PMetal release.

Just a busy week. Will try to look later today!

@nicholasjpaterno
Copy link
Contributor Author

@dcvz appreciate it, let me know if you need any changes to this etc. PMetal depends on a public RandomState API. All the best!

@nicholasjpaterno
Copy link
Contributor Author

@dcvz I definitely understand if its a busy week. Please keep me posted if you think it may be a while before you have a chance to look this over. We missed our launch window on Tuesday but we may move to ship our fork as a temporary work around and deprecate once we can depend on the the upstream fork.

@minghuaw
Copy link
Collaborator

LGTM!

@minghuaw minghuaw self-requested a review February 26, 2026 02:05
@minghuaw minghuaw merged commit d76d388 into oxiglade:main Feb 26, 2026
3 checks passed
@nicholasjpaterno
Copy link
Contributor Author

@minghuaw and @dcvz thanks for taking the time to look into this! I'm updating PMetal now and deprecating my fork 👍 Hoping to ship asap!

@nicholasjpaterno
Copy link
Contributor Author

@minghuaw @dcvz will you be able to ship a minor or patch version release including these changes? Thats the final blocker for PMetal! I have it pinned to the git commit and tested locally, everything is working great!

@nicholasjpaterno
Copy link
Contributor Author

crates.io rejects pinned git dependencies unfortunately

@nicholasjpaterno
Copy link
Contributor Author

nicholasjpaterno commented Feb 26, 2026

I'm happy to do any legwork needed for the release, please let me know. This is the final blocker for PMetal v0.1.0 release

@nicholasjpaterno
Copy link
Contributor Author

Ideally the release would include #320 as well. I can see you're overloaded for now, we will ship our v0.1.0 with our fork for now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants