sharding + vmap

[pharringtonp19]

I have a function which takes a key. I would like to run this function on an array of keys in parallel across a cluster of gpus. Is sharding + vmap the right approach?

import os
import functools
from typing import Optional
import numpy as np
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding


def f(key):
    return jax.random.normal(key)


sharding = PositionalSharding(mesh_utils.create_device_mesh((4,1)))

print(f"\nDefault Device: {jax.devices()[0]}\n")

print("Split Keys")
keys = jax.random.split(jax.random.PRNGKey(0), 12)
jax.debug.visualize_array_sharding(keys)


print("Move Keys")
keys = jax.device_put(keys, sharding)
jax.debug.visualize_array_sharding(keys)


print("Output of Keys")
outputs = jax.vmap(f)(keys)
jax.debug.visualize_array_sharding(outputs)
outputs