I have a batch of `M`

matrices stored as a 3-D array `A`

with shape `(M, N, N)`

and a batch of `M`

vectors stored as a 2-D array `b`

with shape `(M, N)`

. I want to add the `i`

-th vector to the diagonals of the `i`

-th matrix, for all batch indices `i`

. In plain NumPy, this can be done with

```
import numpy as np
A = ... # shape (M, N, N)
b = ... # shape (M, N)
M, N = b.shape
diag_indices = np.arange(N)
A[:, diag_indices, diag_indices] += b
```

I was wondering how to accomplish this in Numba knowing that, from the docs,

A subset of advanced indexing is also supported: only one advanced index is allowed, and it has to be a one-dimensional array (it can be combined with an arbitrary number of basic indices as well).