pmap ValueError

[zohimchandani]

Lets assume I have the following function which I would like to run pmap over:

def some_func(a,b): 
    return a+b 

Given that jax.device_count() returns 2 on my machine, I can run pmap in the following manner for an array of b values where the value for a is broadcasted:

pmap(fun = some_func, backend= 'gpu', in_axes=(None,0))(1, jnp.array([1,2]))

ShardedDeviceArray([2, 3], dtype=int32)

My question is, what if I wanted to run this for more b values than the number of GPUs available? Moreover, I actually want to run this for all combinations for a and b values as shown below:

a = jnp.array([1,2,3,4])
b = jnp.array([1,2,3,4])

pmap(fun = some_func, backend= 'gpu', in_axes=(0,0))(a,b)

[ValueError: compiling computation that requires 4 logical devices, but only 2 XLA devices are available (num_replicas=4, num_partitions=1)]()

where all possible combinations of a and b are executed.

Any ideas on how to approach this problem?

[YouJiacheng]

a = jnp.ones(4)
a = a.reshape((2, 2))
fun = pmap(vmap(lambda a, b: a + b))
fun(a, a)

[zohimchandani]

Your answer works for this small example but doesn't quite help me with my problem.

Let's assume I have this simple function:

def func(a,b): 
    return a + b 

which I want to execute for all possible combinations for a and b and save the result to an empty matrix. In numpy, I would do something like this:

def func(a,b): 
    return a + b 


a = np.arange(0,10)
b = np.arange(0,10)


K = np.zeros(shape = (a.shape[0], a.shape[0]))


for i in range(0, a.shape[0], 1): 
    for j in range(i, a.shape[0], 1): 
        K[i][j] = K[j][i] = func(i,j)

1. How do I do this in Jax given I have 2 GPU devices at my disposal?
2. How do I overcome populating an empty matrix with my answer given that jax arrays are immutable?
3. Can I use pmap with CPU devices?

Thank you.

[YouJiacheng]

Actually you can do in-place update as efficient as numpy within jit, but I don't recommend you to do so.
If func computation is cheap (it is okay to compute (i,j) and (j,i))

import jax
import jax.numpy as jnp

def f(x, y):
  return x + y

fy = jax.vmap(f, in_axes=(None, 0))
fxy = jax.vmap(fy, in_axes=(0, None))

print(fxy(jnp.arange(4), jnp.arange(8)))

fpxy = jax.pmap(fxy, in_axes=(0, None))
print(fpxy(jnp.arange(8).reshape((2, 4)), jnp.arange(8)))

If func computation is costly

import jax
import jax.numpy as jnp

def f(x, y):
  return x + y

def g(x, i, j):
  return f(x[i], x[j])

x = jnp.tile(jnp.arange(4).reshape((4, 1)), 1024)
i, j = jnp.tril_indices(4)
i, j = i.reshape((2, 5)), j.reshape((2, 5))
print(jax.pmap(jax.vmap(g, (None, 0, 0)), in_axes=(None, 0, 0))(x, i, j))

[zohimchandani]

Both methods rely on the fact that given 2 GPUs, one needs to reshape the input args with dimensions (2,x) where x is a whole number. This works when the len of the array to loop over is divisible by 2 without a remainder but not for a general case.

For eg this won't work due to incorrect dimensions:

print(fpxy(jnp.arange(11).reshape((2, ?)), jnp.arange(11)))

How does one in general pmap functions?

[YouJiacheng]

padding to 2x is the only way, IIUC.

[zohimchandani]

Does pmap work with CPUs? I have 20 CPUs at my disposal and 2 GPUs, wondering which one might be faster.