Unable to understand lax.scatter_add and lax.ScatterDimensionNumbers

[ajithmoola]

I have a 2D array PHI of shape (N, 3) and a 1D array of integers num_seg of shape (S, 1) defining a segmentation of PHI, such that jnp.sum(num_seg) would equal to N. The goal is to segment PHI based on num_seg and add up the segmented rows in PHI resulting in a 2D array output of shape (S, 3).

I can easily implement this in PyTorch

operand = torch.zeros((S, 3))
indices = torch.repeat_interleave(torch.arange(N), num_seg).unsqueeze(1).expand(-1, 3)
output = operand.scatter_add_(axis=0, index=indices, PHI)

Similar approach in JAX would look like this

operand = jnp.zeros((S, 3))
indices = jnp.repeat(jnp.arange(N), num_seg).repeat(3, axis=1)
dim_numbers = lax.ScatterDimensionNumbers(scatter_dims_to_operand_dims=(0,), 
                                  update_window_dims=(), 
                                  inserted_window_dims=(0,))
output = lax.scatter_add(operand=operand, scatter_indices=indices, updates=PHI, dimension_numbers=dim_num)

The PyTorch implementation works and yields correct output but the JAX implementation throws an error saying Updates tensor must be of rank 1; got 2.

I figured the issue might be with dim_numbers and went through XLA documentation to read more on lax.ScatterDimensionNumbers but I couldn't understand it. Can anyone throw some light on how lax.ScatterDimensionNumbersworks and help fix the error?

[jakevdp]

My main recommendation here would be to avoid using lax.scatter directly. It is a low-level operation that speaks the language of XLA primitives; generally users will find jnp.ndarray.at a more amenable interface.

For example, if I understand your intent correctly, you could write it like this:

result = operand.at[indices].add(PHI)

If that's not what you have in mind, perhaps you could give a full example of what you're attempting in pytorch (your code above has undefined variables and syntax errors). Thanks!

[ajithmoola]

Thank you so much! Your answer works as I intended.