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

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

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).

Sign up

Join the 1M+ learners on CodeSignal

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

Creating Your First PRNG Key

Let's begin by creating our first PRNG key . We use the jax.random.key() function, providing an integer seed. This function generates what JAX calls a typed PRNG key, that is, a key that explicitly encodes which pseudo-random number generation algorithm it uses, ensuring type safety and allowing JAX to manage different PRNG algorithms efficiently.. JAX's default PRNG algorithm is based on "Threefry," a counter-based generator, which is often abbreviated as "fry" in the key's representation. 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}") When you execute this code, the output will look similar to this: Initial key: Array((), dtype=key<fry>) overlaying: [0 0] Let's break down this output: Array((), dtype=key<fry>) : This tells us key is a JAX array. It's a scalar (empty shape ()), and its dtype (data type) is key<fry> . This special dtype signifies it's a typed PRNG key using the Threefry algorithm. overlaying: [0 0] : This shows a representation of the internal state of the key. For a seed of 0, the internal state is initialized to [0 0] . The integer seed (0 in our example) determines the initial state of the PRNG. Using the same seed will always produce the exact same initial key. This is the starting point for ensuring reproducibility. In practice, you might choose seeds based on experiment IDs, configuration parameters, or other deterministic sources.

Generating Random Numbers and Ensuring Reproducibility

Now that we have a key, let's use it to generate some random numbers. We'll use jax.random.normal() to sample numbers from a standard normal distribution. A key feature of JAX's random system is that using the same key for the same operation will always yield the same result. # 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}") Running this code will produce output like: Random numbers (version 1): [ 1.6226422 2.0252647 -0.43359444] Random numbers (version 2, same key): [ 1.6226422 2.0252647 -0.43359444] As you can see, random_numbers_v1 and random_numbers_v2 are identical! This is because we used the exact same key for both calls to jax.random.normal() with the same shape. This deterministic behavior is fundamental to JAX's design and is what guarantees reproducibility. If you wanted different random numbers, you would need to use a different key for the second call. The general pattern for JAX random functions (like jax.random.normal, jax.random.uniform, etc.) is that the first argument is always the PRNG key. Subsequent arguments specify the parameters of the random operation, such as the shape of the desired output array and any distribution-specific parameters (e.g., minval and maxval for jax.random.randint).

Splitting Keys for Independent Operations

Managing Multiple Random Operations

Now we can use our newly generated subkeys to perform different, independent random operations. Each subkey will fuel one operation, ensuring that the random numbers generated by one operation don't affect the others. # 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}") The output will be similar to this (exact values will depend on the keys generated in the previous step): 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] Here, we've used three different JAX random functions: jax.random.uniform(subkey1, ...): Generates values uniformly distributed in the interval [0, 1). jax.random.normal(subkey2, ...): Generates values from a standard normal distribution. jax.random.randint(subkey3, ...): Generates random integers within a specified range [minval, maxval). Each operation used a distinct subkey (subkey1, subkey2, subkey3). This is the standard practice: one key, one random call. If subkey1 were used again for another call to jax.random.uniform (or any other random function), it would produce the exact same array1 again. To get new random numbers, you'd typically split subkey1 further or use a different subkey. This explicit management makes the flow of randomness clear and prevents the kind of subtle bugs that can arise from shared global random states.

Demonstrating Complete Reproducibility

Let's put everything together to demonstrate how these principles lead to complete reproducibility. We'll simulate two "runs" of a sequence of random operations, both starting from the same initial seed and following the same key-splitting logic. We expect the outputs of corresponding operations in both runs to be identical. 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}") The output will be: 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 Perfect reproducibility! Because we started with the same initial_seed (42) in both runs, key_run1 and key_run2 were identical. Subsequently, splitting these identical keys (jax.random.split(key_run1) and jax.random.split(key_run2)) produced identical pairs of subkeys (key_run1_op1 was identical to key_run2_op1, and key_run1_op2 to key_run2_op2). Finally, using these corresponding identical subkeys in the same random operations (jax.random.normal and jax.random.uniform) yielded identical outputs. This clearly illustrates how JAX's functional approach to randomness, combined with careful key management, enables you to reliably recreate computational results. This is invaluable for debugging, ensuring scientific validity, and consistent model behavior.

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}")

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