Adding a batch of vectors to the corresponding diagonals of a batch of matrices

Numba Support: How do I do ...?
1

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).

2

Hi @christian-cahig,

off the top of my head I cannot provide you with a vectorised expression like in pure numpy. However, when you use numba, the nice thing is that it is perfectly fine to simply use loops. You should not be suffering any performance penalties as you would in CPython.

Cheers

1 Like
3

@Hannes thanks!

I tried reading through this part of docs and came up with this approach:

import numpy as np
import numba as nb

@nb.jit(nopython=True)
def batched_diag_add(A, b):
    M, _ = b.shape
    for m in range(M):
        np.fill_diagonal(A[m], np.diag(A[m]) + b[m])
    ...

This seems to be fine for my use case, but Iā€™m open to other approaches.