Reproducible Randomness with jax.random: Keys, Splitting, and Determinism

Introduction

Welcome to "Advanced JAX: Transformations for Speed & Scale"! I'm thrilled to have you join me for this exciting journey into JAX's more sophisticated capabilities. Having completed "JAX Fundamentals: NumPy Power-Up," you've already mastered the essential building blocks: JAX arrays, pure functions, automatic differentiation with jax.grad, JIT compilation with jax.jit, and structured control flow using jax.lax . Now, we're ready to explore the transformations that make JAX truly powerful for large-scale scientific computing and machine learning. Throughout "Advanced JAX: Transformations for Speed & Scale", we will delve into JAX's explicit random number system (our focus today!), master automatic vectorization with jax.vmap and parallelization across devices with jax.pmap, understand how to work with complex nested data structures using PyTrees, and ultimately apply these skills to build more sophisticated models. In this opening lesson, we'll tackle reproducible randomness using JAX's unique jax.random module. Random number generation is fundamental to many areas, including machine learning — from weight initialization and data shuffling to stochastic optimization and Monte Carlo methods. However, JAX takes a fundamentally different approach to randomness compared to NumPy and most other libraries. Instead of relying on a global, mutable state that can lead to hard-to-debug issues, JAX embraces a functional approach to pseudo-random number generation (PRNG) that ensures complete reproducibility and seamlessly integrates with JAX's transformations. By the end of this lesson, you'll understand how to create and manage PRNG keys, split them for independent operations, and guarantee that your random computations are perfectly reproducible. This foundation will prove invaluable as we progress through the course.

Traditional Stateful Randomness

Traditional libraries like NumPy use stateful random number generation. When you call a function like np.random.normal(), NumPy internally updates a global random state. Each subsequent call then produces the next number in the sequence based on this hidden state. Python import numpy as np # NumPy's stateful approach - each call modifies internal state np.random.seed(42) # Set the global seed print(np.random.normal(size=2)) # First call print(np.random.normal(size=2)) # Second call: different numbers! Notice how each call to np.random.normal() produces different numbers, even though we're not explicitly passing any different parameters. This is because NumPy maintains a hidden global state that gets updated with each random operation. While this stateful approach can be convenient for simple scripts, it has several significant drawbacks, especially in complex or parallel computations: Reproducibility Challenges : It can be difficult to ensure that the exact same sequence of random numbers is generated across different runs or environments, especially if operations are reordered or run in parallel. Parallelism Issues : If multiple parts of your code access the global random state concurrently (e.g., in multi-threaded or distributed settings), they can interfere with each other, leading to non-deterministic behavior. Transformation Incompatibility : JAX's core strength lies in its ability to transform functions (e.g., jit for compilation, vmap for vectorization). A hidden global state doesn't fit well with these functional transformations, which expect functions to be pure (outputs depend only on inputs).

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

import jax
import jax.numpy as jnp

# Generate an initial PRNG key from a seed (e.g., 0)
key = jax.random.key(0)
print(f"Initial key: {key}")

import jax
import jax.numpy as jnp

# Generate an initial PRNG key from a seed (e.g., 0)
key = jax.random.key(0)
print(f"Initial key: {key}")

text

Initial key: Array((), dtype=key<fry>) overlaying:
[0 0]

Initial key: Array((), dtype=key<fry>) overlaying:
[0 0]

Python

# Use the key to sample 3 numbers from a normal distribution
random_numbers_v1 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 1): {random_numbers_v1}")

# Using the exact same key again for the same operation
random_numbers_v2 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 2, same key): {random_numbers_v2}")

# Use the key to sample 3 numbers from a normal distribution
random_numbers_v1 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 1): {random_numbers_v1}")

# Using the exact same key again for the same operation
random_numbers_v2 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 2, same key): {random_numbers_v2}")

text

Random numbers (version 1): [ 1.6226422   2.0252647  -0.43359444]
Random numbers (version 2, same key): [ 1.6226422   2.0252647  -0.43359444]

Random numbers (version 1): [ 1.6226422   2.0252647  -0.43359444]
Random numbers (version 2, same key): [ 1.6226422   2.0252647  -0.43359444]

Python

# Split the key to generate new independent keys.
# The first returned key will update our 'key' variable for future splits.
# subkey1, subkey2, subkey3 are for distinct operations.
key, subkey1, subkey2, subkey3 = jax.random.split(key, num=4)

print(f"New main key after split: {key}")
print(f"Subkey 1: {subkey1}")
print(f"Subkey 2: {subkey2}")
print(f"Subkey 3: {subkey3}")

# Split the key to generate new independent keys.
# The first returned key will update our 'key' variable for future splits.
# subkey1, subkey2, subkey3 are for distinct operations.
key, subkey1, subkey2, subkey3 = jax.random.split(key, num=4)

print(f"New main key after split: {key}")
print(f"Subkey 1: {subkey1}")
print(f"Subkey 2: {subkey2}")
print(f"Subkey 3: {subkey3}")

text

New main key after split: Array((), dtype=key<fry>) overlaying:
[1797259609 2579123966]
Subkey 1: Array((), dtype=key<fry>) overlaying:
[ 928981903 3453687069]
Subkey 2: Array((), dtype=key<fry>) overlaying:
[4146024105 2718843009]
Subkey 3: Array((), dtype=key<fry>) overlaying:
[2467461003 3840466878]

New main key after split: Array((), dtype=key<fry>) overlaying:
[1797259609 2579123966]
Subkey 1: Array((), dtype=key<fry>) overlaying:
[ 928981903 3453687069]
Subkey 2: Array((), dtype=key<fry>) overlaying:
[4146024105 2718843009]
Subkey 3: Array((), dtype=key<fry>) overlaying:
[2467461003 3840466878]

Python

# Use subkey1 to generate a 2x2 array of uniformly distributed numbers
array1 = jax.random.uniform(subkey1, shape=(2, 2))
print(f"Random array 1 (uniform from subkey1):\n{array1}")

# Use subkey2 to generate an array of 2 normally distributed numbers
array2 = jax.random.normal(subkey2, shape=(2,))
print(f"Random array 2 (normal from subkey2):\n{array2}")

# Use subkey3 to generate an array of 5 random integers between 0 and 9
array3 = jax.random.randint(subkey3, shape=(5,), minval=0, maxval=10)
print(f"Random integer array (from subkey3):\n{array3}")

# Use subkey1 to generate a 2x2 array of uniformly distributed numbers
array1 = jax.random.uniform(subkey1, shape=(2, 2))
print(f"Random array 1 (uniform from subkey1):\n{array1}")

# Use subkey2 to generate an array of 2 normally distributed numbers
array2 = jax.random.normal(subkey2, shape=(2,))
print(f"Random array 2 (normal from subkey2):\n{array2}")

# Use subkey3 to generate an array of 5 random integers between 0 and 9
array3 = jax.random.randint(subkey3, shape=(5,), minval=0, maxval=10)
print(f"Random integer array (from subkey3):\n{array3}")

text

Random array 1 (uniform from subkey1):
[[0.00729382 0.02089119]
 [0.5814265  0.36183798]]
Random array 2 (normal from subkey2):
[1.2956359 1.3550105]
Random integer array (from subkey3):
[7 8 1 1 1]

Random array 1 (uniform from subkey1):
[[0.00729382 0.02089119]
 [0.5814265  0.36183798]]
Random array 2 (normal from subkey2):
[1.2956359 1.3550105]
Random integer array (from subkey3):
[7 8 1 1 1]

Python

print("\nDemonstrating reproducibility:")
initial_seed = 42 # A different seed for this demonstration

# --- First run ---
print("--- Run 1 ---")
key_run1 = jax.random.key(initial_seed) # Create key from seed
# Split the key for two operations
key_run1_op1, key_run1_op2 = jax.random.split(key_run1)

output_run1_op1 = jax.random.normal(key_run1_op1, (2,)) # First operation
output_run1_op2 = jax.random.uniform(key_run1_op2, (2,)) # Second operation
print(f"Run 1, Op 1 (normal): {output_run1_op1}")
print(f"Run 1, Op 2 (uniform): {output_run1_op2}")

# --- Second run ---
print("--- Run 2 ---")
key_run2 = jax.random.key(initial_seed) # Re-create key from the SAME seed
# Split the key using the SAME logic
key_run2_op1, key_run2_op2 = jax.random.split(key_run2)

output_run2_op1 = jax.random.normal(key_run2_op1, (2,)) # First operation again
output_run2_op2 = jax.random.uniform(key_run2_op2, (2,)) # Second operation again
print(f"Run 2, Op 1 (normal): {output_run2_op1}")
print(f"Run 2, Op 2 (uniform): {output_run2_op2}")

# Verify that the outputs are identical
are_op1_equal = jnp.array_equal(output_run1_op1, output_run2_op1)
are_op2_equal = jnp.array_equal(output_run1_op2, output_run2_op2)
print(f"Are outputs from Op1 identical across runs? {are_op1_equal}")
print(f"Are outputs from Op2 identical across runs? {are_op2_equal}")

print("\nDemonstrating reproducibility:")
initial_seed = 42 # A different seed for this demonstration

# --- First run ---
print("--- Run 1 ---")
key_run1 = jax.random.key(initial_seed) # Create key from seed
# Split the key for two operations
key_run1_op1, key_run1_op2 = jax.random.split(key_run1)

output_run1_op1 = jax.random.normal(key_run1_op1, (2,)) # First operation
output_run1_op2 = jax.random.uniform(key_run1_op2, (2,)) # Second operation
print(f"Run 1, Op 1 (normal): {output_run1_op1}")
print(f"Run 1, Op 2 (uniform): {output_run1_op2}")

# --- Second run ---
print("--- Run 2 ---")
key_run2 = jax.random.key(initial_seed) # Re-create key from the SAME seed
# Split the key using the SAME logic
key_run2_op1, key_run2_op2 = jax.random.split(key_run2)

output_run2_op1 = jax.random.normal(key_run2_op1, (2,)) # First operation again
output_run2_op2 = jax.random.uniform(key_run2_op2, (2,)) # Second operation again
print(f"Run 2, Op 1 (normal): {output_run2_op1}")
print(f"Run 2, Op 2 (uniform): {output_run2_op2}")

# Verify that the outputs are identical
are_op1_equal = jnp.array_equal(output_run1_op1, output_run2_op1)
are_op2_equal = jnp.array_equal(output_run1_op2, output_run2_op2)
print(f"Are outputs from Op1 identical across runs? {are_op1_equal}")
print(f"Are outputs from Op2 identical across runs? {are_op2_equal}")

text

Demonstrating reproducibility:
--- Run 1 ---
Run 1, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 1, Op 2 (uniform): [0.72766423 0.78786755]
--- Run 2 ---
Run 2, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 2, Op 2 (uniform): [0.72766423 0.78786755]
Are outputs from Op1 identical across runs? True
Are outputs from Op2 identical across runs? True

Demonstrating reproducibility:
--- Run 1 ---
Run 1, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 1, Op 2 (uniform): [0.72766423 0.78786755]
--- Run 2 ---
Run 2, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 2, Op 2 (uniform): [0.72766423 0.78786755]
Are outputs from Op1 identical across runs? True
Are outputs from Op2 identical across runs? True

Python

import jax
import jax.numpy as jnp

# Generate an initial PRNG key from a seed (e.g., 0)
key = jax.random.key(0)
print(f"Initial key: {key}")

text

Initial key: Array((), dtype=key<fry>) overlaying:
[0 0]

Python

# Use the key to sample 3 numbers from a normal distribution
random_numbers_v1 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 1): {random_numbers_v1}")

# Using the exact same key again for the same operation
random_numbers_v2 = jax.random.normal(key, shape=(3,))
print(f"Random numbers (version 2, same key): {random_numbers_v2}")

text

Random numbers (version 1): [ 1.6226422   2.0252647  -0.43359444]
Random numbers (version 2, same key): [ 1.6226422   2.0252647  -0.43359444]

Python

# Split the key to generate new independent keys.
# The first returned key will update our 'key' variable for future splits.
# subkey1, subkey2, subkey3 are for distinct operations.
key, subkey1, subkey2, subkey3 = jax.random.split(key, num=4)

print(f"New main key after split: {key}")
print(f"Subkey 1: {subkey1}")
print(f"Subkey 2: {subkey2}")
print(f"Subkey 3: {subkey3}")

text

New main key after split: Array((), dtype=key<fry>) overlaying:
[1797259609 2579123966]
Subkey 1: Array((), dtype=key<fry>) overlaying:
[ 928981903 3453687069]
Subkey 2: Array((), dtype=key<fry>) overlaying:
[4146024105 2718843009]
Subkey 3: Array((), dtype=key<fry>) overlaying:
[2467461003 3840466878]

Python

# Use subkey1 to generate a 2x2 array of uniformly distributed numbers
array1 = jax.random.uniform(subkey1, shape=(2, 2))
print(f"Random array 1 (uniform from subkey1):\n{array1}")

# Use subkey2 to generate an array of 2 normally distributed numbers
array2 = jax.random.normal(subkey2, shape=(2,))
print(f"Random array 2 (normal from subkey2):\n{array2}")

# Use subkey3 to generate an array of 5 random integers between 0 and 9
array3 = jax.random.randint(subkey3, shape=(5,), minval=0, maxval=10)
print(f"Random integer array (from subkey3):\n{array3}")

text

Random array 1 (uniform from subkey1):
[[0.00729382 0.02089119]
 [0.5814265  0.36183798]]
Random array 2 (normal from subkey2):
[1.2956359 1.3550105]
Random integer array (from subkey3):
[7 8 1 1 1]

Python

print("\nDemonstrating reproducibility:")
initial_seed = 42 # A different seed for this demonstration

# --- First run ---
print("--- Run 1 ---")
key_run1 = jax.random.key(initial_seed) # Create key from seed
# Split the key for two operations
key_run1_op1, key_run1_op2 = jax.random.split(key_run1)

output_run1_op1 = jax.random.normal(key_run1_op1, (2,)) # First operation
output_run1_op2 = jax.random.uniform(key_run1_op2, (2,)) # Second operation
print(f"Run 1, Op 1 (normal): {output_run1_op1}")
print(f"Run 1, Op 2 (uniform): {output_run1_op2}")

# --- Second run ---
print("--- Run 2 ---")
key_run2 = jax.random.key(initial_seed) # Re-create key from the SAME seed
# Split the key using the SAME logic
key_run2_op1, key_run2_op2 = jax.random.split(key_run2)

output_run2_op1 = jax.random.normal(key_run2_op1, (2,)) # First operation again
output_run2_op2 = jax.random.uniform(key_run2_op2, (2,)) # Second operation again
print(f"Run 2, Op 1 (normal): {output_run2_op1}")
print(f"Run 2, Op 2 (uniform): {output_run2_op2}")

# Verify that the outputs are identical
are_op1_equal = jnp.array_equal(output_run1_op1, output_run2_op1)
are_op2_equal = jnp.array_equal(output_run1_op2, output_run2_op2)
print(f"Are outputs from Op1 identical across runs? {are_op1_equal}")
print(f"Are outputs from Op2 identical across runs? {are_op2_equal}")

text

Demonstrating reproducibility:
--- Run 1 ---
Run 1, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 1, Op 2 (uniform): [0.72766423 0.78786755]
--- Run 2 ---
Run 2, Op 1 (normal): [ 0.07592554 -0.48634264]
Run 2, Op 2 (uniform): [0.72766423 0.78786755]
Are outputs from Op1 identical across runs? True
Are outputs from Op2 identical across runs? True