Использование numpy.bmat с numba - PullRequest
Новая
       10

Использование numpy.bmat с numba

1 голос
/ 11 октября 2019

Я пытаюсь использовать np.bmat в моей Python-программе, оптимизированной для Numba. Для этого мне нужно вручную определить сопряженную функцию bmat, поскольку нативная функция от numpy не поддерживается: 

@njit
def _bmat_2d(matrices):
    arr_rows = []
    for row in matrices:
        arr_rows.append(np.concatenate(row, axis=-1))
    return np.array(np.concatenate(arr_rows, axis=0))

(этот код более или менее является упрощенной копией кода из numpy)

Однако:

  1. numba принимает только кортежи на входе np.concatenate [1]
  2. numba очень плохо подходит для преобразования произвольного списка в кортежи [2]

У вас есть идеи для этого?

Ссылки:

1 Ответ

1 голос
/ 11 октября 2019

Подойдет ли следующее для ваших целей? 

import numpy as np
import numba as nb

@nb.njit
def _bmat_2d(m):
    out = np.hstack(m[0])
    for row in m[1:]:
        x = np.hstack(row)
        out = np.vstack((out, x))

    return out

A = np.random.randint(10, size=(3,2))
B = np.random.randint(10, size=(3,1))
C = np.random.randint(10, size=(3,3))
D = np.random.randint(10, size=(4,6))

a = np.bmat(((A, B, C), (D,)))
b = _bmat_2d(((A, B, C), (D,)))

print(np.allclose((a, b))  # True

Обратите внимание, что вы должны передать кортеж кортежей, а не список списков, иначе вы получите "отражение"list "ошибка, поскольку Numba в текущей версии не может обрабатывать списки списков.

...