Как индекс ElementwiseKernel CuPy для нескольких массивов? ⇐ Python

[Anonymous]

Я определяю CuPy ElementwiseKernel для вычислений на графическом процессоре с помощью Python. У меня есть несколько входных и выходных массивов, все многомерные, разной формы, поэтому мне интересно, как назначается переменная i ElementwiseKernel (если бы они все имели форму (m, n), я думаю, что i просто работал бы от 0 до m * n - 1, но здесь у меня много разных m и n).
Чтобы быть конкретным, входы uu, vv и выходы u, v имеют форму (H, W). Входной сигнал cc имеет форму (H, W, 2), а соответствующий выходной сигнал c имеет форму (H, W, 2, 2). Входные данные Ix, Iy, It имеют форму (H + K, W + K) (они получены из «заполнения» данных, из которых получены uu, vv и cc). Наконец, вход R имеет форму (K + 1, K + 1). Остальные пять входных данных являются скалярными.
Природа вычислений является «поэлементной» для uu, vv, cc, что означает, что каждая операция использует одну запись (или, в случае cc, кортеж из двух элементов) в своих вычислениях. Между тем, каждая операция использует небольшое окно записей в Ix, Iy, It и использует каждую запись в R. Мне интересно, например, как из i я мог бы получить индекс строки и столбца моего текущего элемента, скажем, в uu, потому что я не уверен, что i соответствует форме/размеру uu, думаю, это мое текущее предположение, поскольку uu является первым аргументом, и я не могу найти разъяснений в документации.
Я также не уверен, как именно указать "сырой" для них массивы могут повлиять на характер i или индексацию в целом.
Все рекомендации приветствуются.
kernel = cp.ElementwiseKernel(
in_params="raw float32 uu, raw float32 vv, raw float32 cc,"
"raw float32 Ix, raw float32 Iy, raw float32 It,"
"raw float32 R,"
"int32 h, int32 w, int32 ksize, int32 padw, int32 Rw",
out_params="raw float32 u, raw float32 v, raw float32 c",
operation=
"""
int idx = i; // how does this "i" work?
// ... (rest of operation)
"""
)


Подробнее здесь: https://stackoverflow.com/questions/798 ... ple-arrays