NumPy индексы k-й диагонали

Я хотел бы выполнить арифметику с k-й диагональю массива numpy.array. Мне нужны эти индексы. Например, что-то вроде:

>>> a = numpy.eye(2)
>>> a[numpy.diag_indices(a, k=-1)] = 5
>>> a
array([[ 1.,  0.],
       [ 5.,  1.]])

К сожалению, diag_indices возвращает только индексы, составляющие главную диагональ, поэтому в данный момент я делаю:

a += numpy.diag([5], -1)

Но это не кажется таким хорошим или надежным. :-)

Есть ли способ в numpy получить индексы для других, кроме главной диагонали?