simple paralellism with 2 or 4 GPUs

[jecampagne]

Hello,
Looking at the doc on shard_map, I'm quite lost...
If someone can convert this simple example which compute a 2D representation of a f(x,y) function
using either 2 or 4 GPUs it would be help. Here is a 1-GPU (device) code using nested vmap.

import jax
import jax.numpy as jnp
from jax import jit, vmap
from matplotlib import pyplot as plt

@jit
def func(x,y):
    return jnp.sin(jnp.sqrt(x**2+y**2))

Npix=2048
xvals = jnp.linspace(-2*np.pi,2*np.pi,Npix)
yvals = jnp.linspace(-2*np.pi,2*np.pi,Npix)
mapped = vmap(vmap(func, in_axes=(None, 0)), in_axes=(0, None))
result = mapped(xvals,yvals)

plt.figure(figsize=(5,5))
img = plt.imshow(result,cmap=mpl.cm.jet)
plt.show()

I had in mind to split the image in 2x1 or 2x2 grid sending to each device the computation of part of the figure
then glue the results to build the final image.
Thanks for your help.

[jecampagne]

No help?

[jakevdp]

The doc you link to is for more advanced usage – I'd suggest starting with Introduction to sharded computation.

[jecampagne]

Thanks. I will try and come back to share my code...