Numba распараллелить сетку

Question

Я ищу способ переписать код ниже, чтобы numba могла его использовать. В настоящее время при запуске кода выдается ошибка Use of unsupported NumPy function 'numpy.meshgrid' or unsupported use of the function.. По сути, я вычисляю 3-арный катезианский продукт (например, [0 ... 13] x [0..13] x [0..13]).

@numba.jit(nopython=True)
def sequences_of_length(repeat, n_values):
    a = list(range(0, n_values))
    c = [a] * repeat
    x = np.array(np.meshgrid(*c), dtype=np.int16).T.reshape(-1, repeat)
    return x

Например, sequences_of_length(3, 13) должно вывести:

[[ 0  0  0]
 [ 0  1  0]
 [ 0  2  0]
 ...
 [12 10 12]
 [12 11 12]
 [12 12 12]]
(2197, 3)

(где 2197 = 13 ^ 3)

Answer 1

Сегодня я столкнулся с той же проблемой, и мне пришлось выяснить, что, к сожалению, mgrid не входит в число тех функций, которые поддерживает numba. Здесь вы можете увидеть поддерживаемые методы и функции, которые поддерживает numba. Но vstack() поддерживается, и я предполагаю, что следующим шагом вашего mgrid может быть vsatck?