How can I use Numba for Pytorch tensors?
I am new to Numba and I need to use Numba to speed up some Pytorch functions. But I find even a very simple function does not work :(
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)

But the following error appears
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 &lt;class 'torch.Tensor'&gt;
argument 1: cannot determine Numba type of &lt;class 'torch.Tensor'&gt;

Can anyone help me? A link with more examples would be also appreciated. Thanks.

numba supports numpy-arrays but not torch's tensors. There is however a bridge Tensor.numpy():

Returns self tensor as a NumPy ndarray. This tensor and the returned
ndarray share the same underlying storage. Changes to self tensor will
be reflected in the ndarray and vice versa.

That means you have to call jitted functions as:
...
z = vec_add_odd_pos(x.numpy(), y.numpy())

If z should be a torch.Tensor as well, torch.from_numpy is what we need:

Creates a Tensor from a numpy.ndarray.
The returned tensor and ndarray share the same memory. Modifications
to the tensor will be reflected in the ndarray and vice versa. The
returned tensor is not resizable.
...

For our code that means
...
z = torch.from_numpy(vec_add_odd_pos(x.numpy(), y.numpy()))

should be called.
