vmap with reduction

[aespielberg]

I have a situation in my code where I have a very large matrix m, and I want to extract many processed values from it, and stack them in an array. This is a clear-cut application of vmap: I write a function in which I extract values from the array, and I vmap across the indices I want to extract and process, sharing m across all such jobs. This is extremely efficient in JAX.

However, I am a bit stuck on the inverse problem: Given a bunch of jobs that can modify shared matrix m, how can I do this efficiently using vmap? In particular, I am interested in the case where I want to sum all updates to m. In this case, this can be thought of as a map-reduce problem: use a map step to compute all of the modifications that need to be applied to m, then sum them up with m. vmap can obviously perform the first part of the problem, but I am unclear how it can be used for the second part. Is there any way to use vmap here? I cannot return many copies of m and then sum over them, since m can be very large. I was hoping there might be something like a vreduce.

I saw two other functions jax.lax.scatter and jax.lax.scatter_add. These look like they may be the functions that I need. However, do they parallelize efficiently on the gpu like vmap? Also, are there any simple examples with them? The documentation seems to be lacking.

EDIT:
Here is a short example, from my code (EDIT 2: Now runs self-contained). The goal here is to construct matrices grid_m_in and grid_v_in. This is done by looping over a very large number, n_particles, computing two updates ( a * mask and jnp.squeeze(weight * mass * mask)), and then accumulating them to grid_m_in and grid_v_in. This should be almost embarrassingly parallel, quite sparsely parallelizable for the computation of the update, just not parallelizable for the accumulation/reduction. I have tried brute-force vmapping and then summing the output matrices, but this is much slower than I think it should be (orders of magnitude slower than extracting values from similar matrices) and will use much more memory.

import jax.numpy as jnp
from jax import grad, jit, vmap
import math
import numpy as np
import jax.nn as jnn
import timeit
import jax

n_particles = 3000
n_grid = 128

dx = 1./n_grid

inv_dx = 1./n_grid

dim = 2
mass = 1.0


def p2g_accum(x, v, affine):
  grid_v_in = jnp.zeros((n_grid, n_grid, dim))
  
  #compute some indexing information
  base = jnp.int32(x * inv_dx - 0.5)
  fx = x * inv_dx - base
  w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
  
  for i in range(3):
    for j in range(3):
      #compute some more indexing information
      offset = jnp.array([i, j])
      dpos = (jnp.array([i, j]) - fx) * dx
      weight = w[i][0] * w[j][1]
      idx = base + offset
      
      #make masks so we don't have to perform slow indexing
      mask = jnn.one_hot(idx[0] * n_grid + idx[1], n_grid * n_grid)
      mask = jnp.expand_dims(jnp.reshape(mask, (n_grid, n_grid)), -1)
      
      #now let's do some complex computation
      a = jnp.expand_dims(jnp.expand_dims(weight * (mass * v + dpos @ affine), 0), 0)
      grid_v_in +=  a * mask
  
  return grid_v_in
  
p2_j = vmap(p2g_accum)

def p2g_reduce(x, v, affine):
  grid_v_in = p2_j(x, v, affine)
  grid_v_in = jnp.sum(grid_v_in, axis=0)
  return grid_v_in
  
  
p2r = jax.checkpoint(jit(p2g_reduce))
  
if __name__ == '__main__':
  
  x = jnp.array(np.random.rand(n_particles, dim))
  v = jnp.array(np.random.rand(n_particles, dim))
  affine = jnp.array(np.random.rand(n_particles, dim, dim))
  p2r(x, v, affine)
  number = 10
  print(timeit.timeit(lambda : p2r(x, v, affine).block_until_ready(), number=number) / number)

[jakevdp]

Hi - thanks for the question. I'm a bit unclear on what kind of operation you're describing. Can you edit the question with a short example of the kind of operation you want to do? What do you mean by "a bunch of jobs that can modify shared matrix m"?

[jakevdp]

Looking at your example, I suspect the real issue is you're trying to use dense matrices for a sparse algorithm. Your arrays have an 0.05% average fill factor:

out = p2_j(x, v, affine)
print(jnp.sum(out != 0) / out.size)
# 0.0005493164

If I were you, I'd stop worrying about how to optimize loops vs vmaps, and start looking into sparse representations. You could start with some of the tools here: https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html

[aespielberg]

I do plan to look into sparse arrays, but would like to avoid the experimental branch for now. However, shouldn't it be possible to rewrite this code like this? Assuming it's efficient, he only issue here is how to use the scatter_add, for which I can find zero examples of how to use online.

import jax.numpy as jnp
from jax import grad, jit, vmap
import math
import numpy as np
import jax.nn as jnn
import timeit
import jax

n_particles = 3000
n_grid = 128

dx = 1./n_grid

inv_dx = 1./n_grid

dim = 2
mass = 1.0

kernel_size = 3

index_array = np.zeros((kernel_size, kernel_size, dim))
  
for i in range(kernel_size):
  for j in range(kernel_size):
    index_array[i, j] = np.array([i, j])

def p2g_accum_inner(x, v, affine, w, fx, base, index):
  offset = index
  dpos = (index - fx) * dx
  weight = w[i][0] * w[j][1]
  idx = base + offset
  
  #now let's do some complex computation
  grid_v_in =  mass * v + dpos @ affine
  return grid_v_in, idx
  
axes = (None, None, None, None, None, None, 0)
p2gai = vmap(vmap(p2g_accum_inner, in_axes=axes), in_axes=axes)
  

def p2g_accum(x, v, affine):
  
  
  #compute some indexing information
  base = jnp.int32(x * inv_dx - 0.5)
  fx = x * inv_dx - base
  w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
  
  grid_v_in, idx = p2gai(x, v, affine, w, fx, base, index_array)
  
  return grid_v_in, idx
  
p2_j = vmap(p2g_accum)

def p2g_reduce(x, v, affine):
  grid_v_in = jnp.zeros((n_grid, n_grid, dim))
  grid_v_in_scatter, idx_scatter = p2_j(x, v, affine)
  grid_v_in = jax.lax.scatter_add(grid_v_in, idx_scatter, grid_v_in_scatter, jnp.array((3, 3, 2))) #THIS LINE IS WRONG
  return grid_v_in
  
  
p2r = jax.checkpoint(jit(p2g_reduce))
  
if __name__ == '__main__':
  
  x = jnp.array(np.random.rand(n_particles, dim))
  v = jnp.array(np.random.rand(n_particles, dim))
  affine = jnp.array(np.random.rand(n_particles, dim, dim))
  p2r(x, v, affine)
  number = 10
  print(timeit.timeit(lambda : p2r(x, v, affine).block_until_ready(), number=number) / number)

[jakevdp]

Here's an approach that I think matches your example, but rather than returning a sparse array to reduce, it returns indices and values to reduce. This should be much faster than your original approach:

def p2g_accum(x, v, affine):
  #compute some indexing information
  base = jnp.int32(x * inv_dx - 0.5)
  fx = x * inv_dx - base
  w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]

  idx = jnp.zeros((3, 3, 2), dtype=int)
  vals = jnp.zeros((3, 3, 2), dtype=float)
  
  for i in range(3):
    for j in range(3):
      offset = jnp.array([i, j])
      dpos = (jnp.array([i, j]) - fx) * dx
      weight = w[i][0] * w[j][1]
      idx = idx.at[i, j].set(base + offset)
      vals = vals.at[i, j].set(weight * (mass * v + dpos @ affine))
  
  return idx, vals

idx, vals = vmap(p2g_accum)(x, v, affine)
result = jnp.zeros((n_grid, n_grid, dim)).at[idx[..., 0], idx[..., 1]].add(vals)

[aespielberg]

Ohhh, that is EXACTLY what I was trying to do, and it gives toughly the performance I expected. I had no idea about the at/add syntax. I thought at indexing was slow in my experience; I'm guessing though, that if performed in batch, it is quite fast?

[jakevdp]

at indexing is slow-ish outside JIT compared to the equivalent numpy operation. Inside JIT it is often fused with other operations so that the cost becomes negligible.

For more on this scatter-update syntax, see https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at