What is the correct approach for generating pseudo random numbers in parallel using sharding in jax? The following doesn't work (due to sampling the same chain) sharding = jax.sharding.PositionalSharding( jax.experimental.mesh_utils.create_device_mesh((8,)), ) @jax.jit(static_argnum=1, out_sharding=sharding.reshape(8, 1)) def uniform_sharded(rng_key, n): return jax.random.uniform(key=rng_key, shape=(n,)) I considered performing a pmap across devices with an array of keys, but the output would then be dependent on the device count and would seem to defeat the intent of sharding. Continue reading...
