小编roz*_*ang的帖子

我是 Numba 新手，我需要使用 Numba 来加速一些 Pytorch 功能。但我发现即使是一个非常简单的功能也不起作用:(

import torch
import numba

@numba.njit()
def vec_add_odd_pos(a, b):
    res = 0.
    for pos in range(len(a)):
        if pos % 2 == 0:
            res += a[pos] + b[pos]
    return res

x = torch.tensor([3, 4, 5.])
y = torch.tensor([-2, 0, 1.])
z = vec_add_odd_pos(x, y)

但出现以下错误

def vec_add_odd_pos(a, b):
    res = 0.
    ^

This error may have been caused by the following argument(s):
- argument 0: cannot determine Numba type of <class 'torch.Tensor'>
- argument 1: …

例如，如果我有

vec a(3, fill::randu);
vec b(5, fill::randu);

如何获得大小为 8 的新向量 c，其中前三个元素来自 a，其余元素来自 b？

我正在使用PyTorch多处理进行编程。我希望所有子过程都可以读/写相同的张量列表（不调整大小）。例如，变量可以是

m = list(torch.randn(3), torch.randn(5))

由于每个张量具有不同的大小，因此我无法将它们组织成一个张量。

python列表没有share_memory_（）函数，并且multiprocessing.Manager无法处理张量列表。如何在多个子流程之间共享变量m？

我想编写一个在 CPU 模式和 CUDA 设备模式下使用的 Numba 函数。当然，我可以使用和不使用 cuda.jit 装饰器编写两个相同的函数。例如：

from numba import cuda, njit

@njit("i4(i4, i4)")
def func_cpu(a, b)
    return a + b

@cuda.jit("i4(i4, i4)", device=True)
def func_gpu(a, b)
    return a + b

但在软件工程中它是丑陋的。有没有一种更优雅的方式，即将代码组合在一个函数中？