Seeking simple example of jax.lax.scatter_add

[olumide-x]

I'm struggling to understand the arguments to jax.lax.scatter_add. I've read the docs but it doesn't help much.

I'd appreciate a very simple example of the jax.lax.scatter_add function.

Thanks.

[hawkinsp]

I strongly recommend you use the x.at[idx].add(y) syntax instead!

The XLA scatter op is notoriously confusing, but for most use cases you should use the NumPy-style indexed update syntax extension that JAX offers. Does that work for you?

[soraros]

@olumide-x IMO, you should think in terms of vmaps instead. Of course, there are many ways to implement such a thing, following is just one attempt

from itertools import product

import numpy as np
import jax.numpy as jnp
from jax import grad, vmap

N = 3
x = y = jnp.arange(N**2.).reshape(N, N)

def W(*args):
  assert all(a.ndim == 0 for a in args)
  return sum(args)

def pieces(x):
  return [x[i] for i in product([np.s_[:-1], np.s_[1:]], repeat=2)]

def total_energy(x, y):
  return jnp.sum(vmap(vmap(W, in_axes=0), in_axes=0)(*pieces(x), *pieces(y)))

grad(total_energy, argnums=(0, 1))(x, y)

[olumide-x]

@soraros I'm really struggling to understand why the nested vmap works.

[soraros]

@olumide-x What your SO question asked for is really mapping W over (the first 2 dims, hence the nested vmap, of) sub-domains of the x and y fields simultaneously and then summing up the results.

[olumide-x]

@soraros Sadly, despite a great deal of deliberation, I have to admit defeat and confess that vmap(vmap(W, in_axes=0), in_axes=0) still eludes me.

Can you please break down what is happening here? Why is the energy function applied twice -- if that's what is happening. That's what it looks like to me. I expect the function is applied just once for each element. Each invocation of W is passed the eight coordinates that define an element, specifically:

W(x0, y0, x1, y1, x4, y4, x3, y3)
W(x1, y1, x2, y2, x5, y5, x4, y4)
W(x4, y4, x5, y5, x8, y8, x7, y7)
W(x3, y3, x4, y4, x7, y7, x6, y6)

Yet in your example you appear to be passing all the coordinates to W. Maybe I don't understand vmap well enough. I struggle a lot with vmap.

As an aside, I did a make_jaxpr on your example and, reassuringly, I can see lots of scatter-adds -- which appears to reinforce my earlier intuition. However the blocks of instructions in the jaxpr are repetitive, even after a jit. Does XLA compile optimize the repetition? I am concerned about the implications for instruction caches especially for large problems as well as the pattern of memory access. On the second point, would XLA attempt to initialze the Jacobian array in a cache friendly/cache-efficient manner?

[olumide-x]

@soraros Finally. I understand your answer. Its only taken a week and a half to do so 😆 Thanks

[olumide-x]

Although I am still trying to understand @soraros's reply (the nested vmap is a source of much consternation to me), I have figured out the arguments to lax.scatter_add, albeit using (the debugger) PDB on the repo https://github.com/tristandeleu/metax. Here's how:

Clone the repo and cd to its root. In the REPL:

>> import pdb
>> pdb.runcall(metax.tests.test_prototypes.test_get_prototypes)

Break on scatter_add and you will see its arguments

operand = DeviceArray([[0., 0., 0., 0., 0., 0., 0.],
             [0., 0., 0., 0., 0., 0., 0.],
             [0., 0., 0., 0., 0., 0., 0.]], dtype=float32)
scatter_indices = DeviceArray([[1],
             [0],
             [2],
             [1],
             [0]], dtype=int32)
updates = DeviceArray([[-0.19596423,  0.254582  ,  0.573146  ,  0.44064867,
              -0.847186  ,  0.31778416,  0.64643025],
             [ 0.03368271, -0.88885146, -0.26579142, -1.5609733 ,
              -0.63806945, -0.44171792,  0.9098043 ],
             [-0.01651609,  0.7582043 ,  1.0892068 ,  0.29885298,
               1.490981  ,  0.07877276, -1.222362  ],
             [ 0.951683  ,  0.21038245,  1.3863313 , -0.33807516,
               2.9521072 , -0.99476653, -0.515637  ],
             [ 0.2918976 , -0.14347766,  1.614347  ,  1.6433928 ,
               0.11422886,  0.25447085, -1.3060236 ]], dtype=float32)
dimension_numbers = ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,))
indices_are_sorted = False
unique_indices = False
mode = None

Return value

DeviceArray([[ 0.3255803 , -1.0323291 ,  1.3485556 ,  0.08241951,
              -0.5238406 , -0.18724707, -0.3962193 ],
             [ 0.75571877,  0.46496445,  1.9594773 ,  0.10257351,
               2.104921  , -0.6769824 ,  0.13079327],
             [-0.01651609,  0.7582043 ,  1.0892068 ,  0.29885298,
               1.490981  ,  0.07877276, -1.222362  ]], dtype=float32)