scatter and update

[AakashKumarNain]

I have a jax array of shape (m, n, ...). Given the dim number, I wanted to update the entries of this array from another array. For example,

source = [[77, 82, 29, 64, 62, 33, 97], [46, 50,  9, 54, 66, 58, 82]]
target = [[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0]] # any array of same shape as source
indices = [[0, 2], [1, 4]]

dim = 1 # If this is 1, then indices represents column numbers for each row
output: [[77,  0, 82,  0,  0,  0,  0], [ 0, 46,  0,  0, 50,  0,  0]]

I am not sure how to do this with at in a nice way

[jakevdp]

I suspect the most convenient way to do this will be using ndarray.at with vmap:

import jax
import jax.numpy as jnp

source = jnp.array([[77, 82, 29, 64, 62, 33, 97], [46, 50,  9, 54, 66, 58, 82]])
target = jnp.array([[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0]])
indices = jnp.array([[0, 2], [1, 4]])

@jax.vmap
def f(source, target, indices):
  return target.at[indices].set(source[:len(indices)])

print(f(source, target, indices))
# [[77  0 82  0  0  0  0]
#  [ 0 46  0  0 50  0  0]]

[AakashKumarNain]

Thanks @jakevdp. This works!