问使用numba.guvectorize的范围的索引

EN

def compute_kinetic_energy(velocity, mass, ke):
    ke = 0.0

    # Increase kinetic energy
    for i in prange(velocity.shape[0]):
        for x in range(2):
            ke += velocity[i, x] * velocity[i, x] * mass[i]

compute_kinetic_energy_gu = guvectorize(['float32[:,:], float32[:], float32'],
                             '(nnodes,dim),(nnodes),()->()',
                             target='cpu',
                             nopython=True)(compute_kinetic_energy)

import numpy as np
from numba import guvectorize, prange

nnodes = 1000
mass = np.ones(nnodes, dtype=np.float32)
velocity = np.zeros(nnodes*3, dtype=np.float32).reshape(nnodes,3)

compute_kinetic_energy(velocity, mass, ke)      # Works :)
compute_kinetic_energy_gu(velocity, mass, ke)   # Do not work :(