PRNGKey handling with sharding + jit

[giovannicemin]

Hello everyone,

I recently ran into some issues handling random keys within functions with sharded inputs.The main challenge is that jax.random.split only accepts single keys, while sharding passes batched data to the function. I came up with a workaround that I'd like to share:

If this solution works for you, great!
If you have a better approach, I'd love to hear your suggestions.

A minimal example is:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
from functools import partial
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

def random_choice(key, a):
    return jax.random.choice(key, a)

def f(key, a, b):
    k1, k2 = jax.random.split(key)
    vmap_rc = jax.vmap(random_choice, in_axes=[0,None])
    return (vmap_rc(k1, a), vmap_rc(k2,b))

mesh = Mesh(jax.devices(), axis_names=('i',))
sharding = NamedSharding(mesh, P('i'))
replicated_sharding = NamedSharding(mesh, P())

key = jax.random.key(0)
keys = jax.random.split(key, 8)
keys = jax.device_put(keys, sharding)

array = jax.device_put(jnp.arange(10), replicated_sharding)
array_ = jax.device_put(jnp.arange(10), replicated_sharding)

jax.jit(f)(keys, array, array_)

This throws the following error: ValueError: split accepts a single key, but was given a key array of shape (8,) != (). Use jax.vmap for batching.

The solution (or workaround) that I found is to replace jax.random.split with:

def batched_key_split(key: chex.PRNGKey, num: int = 2):
    return jax.vmap(jax.random.split, in_axes=[0, None])(key, num).T

in this way, the above code spits out:

(Array([1, 6, 9, 9, 8, 2, 1, 6], dtype=int32),
 Array([0, 6, 2, 5, 4, 1, 2, 7], dtype=int32))

[jakevdp]

Wrapping split with vmap certainly works – but since you're now vmapping every individual operation within f, why not just vmap the whole function?

@partial(jax.vmap, in_axes=(0, None, None))
def f(key, a, b):
    k1, k2 = jax.random.split(key)
    return random_choice(k1, a), random_choice(k2, b)

[giovannicemin]

Hi @jakevdp

Thank you for the answer; indeed you are right. This example helped me realize that vmap is more versatile than I initially thought—I can apply it at nearly any level. However, in my specific implementation, applying vmap to the entire function ended up slowing down the code.
In the end, I choose to use shard_map, which gives a more precise handling of the input/output shards, where the spec for the key is P(). This way I can split the keys normally and then I vmap only the functions that I need to.

Thank you for your response; you’re absolutely right. This example helped me realize that vmap is more versatile than I initially thought—I can apply it at nearly any level. However, in my specific implementation, applying vmap to the entire function ended up slowing down the code.

Ultimately, I decided to use shard_map instead, as it offers a precise control over input and output shards. For example, I can specify the key with P(), then I can split the keys normally, and apply vmap only to the functions that benefit from it.