numpy array: fast assign short array to large array with index

Question

I want to assign values to large array from short arrays with indexing. Simple codes are as follows:

import numpy as np

def assign_x():
    a = np.zeros((int(3e6), 20))
    index_a = np.random.randint(int(3e6), size=(int(3e6), 20))
    b = np.random.randn(1000, 20)
    for i in range(20):
        index_b = np.random.randint(1000, size=int(3e6))
        a[index_a[:, i], i] = b[index_b, i]
    return a

%timeit x = assign_x()
# 2.79 s ± 18.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I have tried other ways which may be relevant, for example np.take and numba's jit, but seems the above is the fastest way. It can also be possibly speed-up using multiprocessing. I have profile the codes, the most time is at the line below as it runs many times (20 here)

a[index_a[:, i], i] = b[index_b, i]

Any chance I can make this faster before using multiprocessing?