为何在Numba中无法执行Numpy数组索引赋值操作？

代码示例 #

import numpy as np    
import matplotlib.pyplot as plt  
import random  
import time  
import math  
import numba  
from numba import jit,prange  

nx=100  
p1=np.zeros([nx],dtype=float)  

@jit(nopython=True)  
def f(r):  
    p1[2]=10  
    return p1[2]   

u3=f(2)  

报错信息 #

Compilation is falling back to object mode WITH looplifting enabled because Function "f" failed type inference due to: No implementation of function Function() found for signature:

setitem(readonly array(float64, 1d, C), Literalint, Literalint)

There are 16 candidate implementations:
- Of which 14 did not match due to:
Overload of function 'setitem': File: : Line N/A.
With argument(s): '(readonly array(float64, 1d, C), int64, int64)':
No match.
- Of which 2 did not match due to:
Overload in function 'SetItemBuffer.generic': File: numba\core\typing\arraydecl.py: Line 176.
With argument(s): '(readonly array(float64, 1d, C), int64, int64)':
Rejected as the implementation raised a specific error:
NumbaTypeError: Cannot modify readonly array of type: readonly array(float64, 1d, C)

观察现象 #

• 不使用索引赋值（如y[i]=x[i]**2）时，数组在函数内可正常工作
• 直接执行数组整体运算（如y=x**2）可行

原因 #

Numba的nopython模式会把全局作用域的Numpy数组标记为只读对象，禁止在JIT编译的函数内直接修改其元素。而数组整体运算（如x**2）是创建新数组，不会改动原全局数组，因此能正常执行。

解决方法 #

1. 将数组作为函数参数传入

把需要修改的数组作为参数传递给JIT函数，Numba会将其识别为可修改的输入：

@jit(nopython=True)  
def f(r, arr):  
    arr[2] = 10  
    return arr[2]   

u3 = f(2, p1)

2. 使用Numba typed List（适合全局可修改场景）

如果必须使用全局可修改的数组结构，可以用Numba提供的typed.List：

from numba.typed import List

p1 = List()
p1.extend(np.zeros(nx, dtype=float))

@jit(nopython=True)  
def f(r):  
    p1[2] = 10  
    return p1[2]   

u3 = f(2)

注意点 #

• 避免在JIT函数中依赖全局变量，不仅会引发只读问题，还会降低代码可维护性和并行安全性
• 数组整体运算返回的是新数组，原全局数组不会被修改，和索引赋值的行为完全不同

内容的提问来源于stack exchange，提问作者AB1998