Let's say I have an array a of shape (3, 3, 3) and a vector b of size 3. I want to iterate on the second axis of the first array, and multiply the first array by the first term of b, the second array by the second term of b etc.

In python I would simply write something like

a = np.ones([3, 3, 3])
b = np.array([1, 2, 3])
for i in range(3):
  a[:, i, :] *= b[i]

(I could even use vectorized functions to do this).

I'm strugling to do it in a clean way in js, as I don't want to loop on every axis (I have matrices with 6 axes, it would be very unefficient.

Currently I use NumJs, and I do it this way:

var a = nj.ones([3, 3, 3]);
let b = [1, 2, 3];
for (let i = 0; i < 3; i++) {
  for (let j = 0; j < 3; j++) {
    a.slice([i, i+1], [j, j+1]).multiply(b[j], false)
  }
}

I have to loop on 2 axes and it doesn't look very efficient.

Are there better ways to do this kind of operations in js? (I would also like to do other operations like to sum an array on some axes only) Thanks!