https://github.com/yichengdwu/random123

Splittable pseudorandom number generators
https://github.com/yichengdwu/random123

mojo

Last synced: 4 months ago
JSON representation

Splittable pseudorandom number generators

• Host: GitHub
• URL: https://github.com/yichengdwu/random123
• Owner: YichengDWu
• License: mit
• Created: 2024-06-19T02:51:52.000Z (about 2 years ago)
• Default Branch: main
• Last Pushed: 2024-06-21T02:29:49.000Z (about 2 years ago)
• Last Synced: 2026-01-13T06:11:59.095Z (6 months ago)
• Topics: mojo
• Language: Mojo
• Homepage:
• Size: 85 KB
• Stars: 1
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

          
# Random123

Random123 is a library of *splittable* pseudo random number generators. It follows [Jax's PRGN design](https://jax.readthedocs.io/en/latest/jep/263-prng.html).

# Usage

```mojo

import random123

from random123 import PRNGKey

var key = random123.key(123)

var samples = random123.normal[DType.float32](key^, 1000)

fn f(owned key: PRNGKey):

    var newkey: PRNGKey

    var subkey: PRNGKey

    newkey, subkey = random123.split(key^, 2)

    g(newkey^)

    h(subkey^)

```

Implemented methods include `bits`, `uniform` and `normal`. More to come!

Note that it is required to use the transfer operator `^` whenever using a `PRNGKey`, and `split` it into new keys if needed.

Otherwise, you should make sure every key is only used once.

# Benchmarks

Random123 is blazingly fast, thanks to the buildin SIMD vectors and `parallelize` in Mojo.

Here is a comparison with Jax:

![Benchmark Results](./benchmark_chart.png)

# Heads-up

You might bump into a compiler bug with `DType.float16` or `DType.bfloat16`. 

See the [issue](https://github.com/modularml/mojo/issues/3073#issue).