如何修改Python函数使其能够使用Numba，以便可以在Google Colab GPU / TPU上加快速度？

Question

我有一个很慢的python函数，我相信如果在GPU / TPU上处理它会更快。我正在使用Google Colab。如何修改它，以便可以使用Numba在Google Colab GPU上对其进行处理？

def Tournament (x, w, cutoff, y):
  preds = np.matmul(x,w)
  preds2 = (preds<np.quantile(preds, 1-cutoff,axis=0))
  preds = (preds>np.quantile(preds, cutoff,axis=0))
  rets = y * preds
  rets[rets == 0] = np.nan
  rets2 = y * preds2
  rets2[rets == 0] = np.nan
  ans = np.nanmedian(rets, axis=0) - np.nanmedian(rets2,axis=0)
  ans = ans[:,None]
  rand2 = np.random.uniform(-0.00000001,0.00000001,size=(popsize,1))
  ans += rand2
  garbage,sort = np.unique(ans,return_index = True)
  sort = sort[:,None]
  return sort, ans

我已经尝试加载Numba库，并在该函数之前使用@jit，但似乎不起作用。