Add support for np.diag_indices() and np.diag_indices_from()

[KrisMinchev]

This PR improves on #4074 by implementing the functions np.diag_indices() and np.diag_indices_from(). Note that the NumPy versions return tuples whereas the implementations presented here return NumPy arrays.

[guilhermeleobas]

Hi @KrisMinchev, great work on adding these functions to Numba.

Note that the NumPy versions return tuples whereas the implementations presented here return NumPy arrays.

Numba tries to keep compatibility with NumPy as close as possible. I've left a suggestion on how you can return tuples instead of NumPy arrays. The idea is to create a tuple at compile-time with the same type as the tuple you will return at runtime, and use tuple_setitem to set each tuple item.

[guilhermeleobas]

Numba has a function called tuple_setitem from cpython/unsafe/tuple.py where you can use to set items in a tuple at runtime.

Suggested change

	def impl(n, ndim=2):
	res = np.arange(n * ndim)
	for i in range(n * ndim):
	res[i] = res[i] % n
	return res.reshape((ndim, n))
	# tup_init is used only to get the correct tuple type
	tup_init = (np.arange(1),) * 2
	from numba.cpython.unsafe.tuple import tuple_setitem
	def impl(n, ndim=2):
	res = np.arange(n * ndim)
	for i in range(n * ndim):
	res[i] = res[i] % n
	x = res.reshape((ndim, n))
	tup = tup_init
	for i in range(n):
	tup = tuple_setitem(tup, i, x[i])
	return tup

[KrisMinchev]

Hi @guilhermeleobas, thank you for your suggestion. The only issue with this implementation is that the array elements that are being returned are immutable, which is different to what the NumPy function np.diagflat returns:

import numpy as np
from numba import njit
@njit
def foodiagindices(n, ndim=2):
    res = np.diag_indices(n, ndim)
    return res
a=np.diag_indices(5)
print(a)
a[0][0] = 5
print(a)
c=foodiagindices(5)
print(c)
c[0][0] = 5

Returns

(array([0, 1, 2, 3, 4]), array([0, 1, 2, 3, 4]))
(array([5, 1, 2, 3, 4]), array([5, 1, 2, 3, 4]))
(array([0, 1, 2, 3, 4]), array([0, 1, 2, 3, 4]))
Traceback (most recent call last):
File "/home/kristian/Desktop/numba/a.py", line 13, in
c[0][0] = 5
ValueError: assignment destination is read-only

Another issue is that your proposed implementation implicitly assumes that ndim is equal to 2 at runtime by setting tup_init = (np.arange(1),) * 2. Hence, the following code:

import numpy as np
from numba import njit
@njit
def foodiagindices(n, ndim=2):
    res = np.diag_indices(n, ndim)
    return res
c = foodiagindices(5, 3)

throws a segmentation fault.

[guilhermeleobas]

Another issue is that your proposed implementation implicitly assumes that ndim is equal to 2 at runtime by setting tup_init = (np.arange(1),) * 2. Hence, the following code

Oh, I know. My code was just an example on how to use tuple_setitem

[KrisMinchev]

The current implementation does not work as intended. It is hard to keep the behaviour of the NumPy function np.diag_indices(n, ndim) since the size of the tuple being is equal to ndim and hence cannot be determined at compile time. Since np.diag_indices_from() calls np.diag_indices(), the same issue is present there as well.

[guilhermeleobas]

In today's public meeting (notes), it was decided to support diag_indices with ndim as a literal argument.

Kristian, do you have any other question regarding supporting this function?