Using psum inside a function passed to jax.scipy.sparse.linalg.cg

[maurorigo]

I have to solve a linear system and, due to memory constraints, I'm currently using jax.scipy.sparse.linalg.cg. To speed things up and save memory, I'm using multiple GPUs. The code of a pmapped function (with pmapped axis 'p') looks something like this:

def pmapped_func(self, pmap_inputs):
        pmap_var1 = some_func(pmap_inputs)
        pmap_var1 = pmap_var1 - psum(jnp.sum(pmap_var1, axis=0), axis_name='p')
        (...)
        f = lambda v: psum(jnp.matmul(jnp.matmul(pmap_var1, v), pmap_var1), axis_name='p')
        out, info = jax.scipy.sparse.linalg.cg(f, var2, maxiter=1000)

        return out

Is having a psum inside cg a sensible thing to do? I'm asking this because it seems like the more devices I use, the longer it actually takes to solve the system (with more than 2 GPUs). Could this be due to the communications between the GPUs begin relatively slow?

Also, cg is inside the pmapped function because the full version of pmap_var1 (not distributed across the GPUs) is large, so having it divided between the GPUs is more convenient, but at the same time each GPU calls cg; is there a smart way to make it such that only one GPU calls it (without using lax.cond inside the pmapped function for example)?