使用带有 for 循环的 jited 函数

Question

我有以下代码

@njit
    def ss(w,g,sm):
        kb = ( (α/(r*(1+τk)))**((1-γ)/(1-γ - α)) )* \
        ( (γ/(w*(1+τn)))**(γ/(1-γ-α)) )* \
        (smat*(1-τo))**(1/(1-γ-α)) ## ss capital

        nb = (1+τk)*r*γ/((1+τn)*w*α)*kb ## ss labor

        πb = (1-τo)*sm*(kb**α)*(nb**γ)- (1+τn)*w*nb-(1+τk)*r*kb-cf #ss profit
        W = πb/(1-0.0196) ## error in the code
        for i in range(ns):
            for j in range(nτ):
                xb[i,j] = 1 if W[i,j]>=0 else xb[i,j]
        we = sum(W*g*xb) - ce
        return we

据我所知它应该可以工作，但我不断收到以下错误

TypingError: Invalid use of Function(<built-in function setitem>) with argument(s) of type(s): (readonly array(float64, 2d, C), tuple(int64 x 2), float64)
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:

我仍在寻找解决 python 的方法。我知道这是因为我正在使用 njit，但究竟是什么部分导致了它？如果我删除 for 循环，它可以正常工作，但我想知道 for 循环中的什么可能导致这种情况？

可重现的例子

import numpy as np
from numba import njit
class rep:
    def __init__(self, A = 1, B=2, C = 3, D = 4):
        self.A, self.B,self.C, self.D = A,B,C,D


def op(cls):
    A,B,C,D = cls.A, cls.B,cls.C, cls.D
    xb = np.zeros([10,10])
    @njit
    def rep1(w,g,sm):
        kb = A*3 + B*2
        nb = C*3 - D*w
        pib = sm*kb - nb
        W = pib/4
        for i in range(10):
            for j in range(10):
                xb[i,j] = 1 if W[i,j]>=0 else xb[i,j]
        we = np.sum(w*xb*g)
        return we
    return rep1


g = np.zeros([10,10])
sm = np.ones([10,10])
w = 1
r = rep()
rep1 = op(r)
print(rep1(w,g,sm))

Answer 1

这似乎与 cmets 中提到的 @max9111 相关 - xb 数组是在 jited 函数外部创建的，因此无法在内部修改。我通过将xb 移出op 并将其作为参数传递给rep1 对您的代码进行了一些更改，并且可以成功执行：

import numpy as np
from numba import njit
class rep:
    def __init__(self, A = 1, B=2, C = 3, D = 4):
        self.A, self.B,self.C, self.D = A,B,C,D


def op(cls):
    A,B,C,D = cls.A, cls.B,cls.C, cls.D
    @njit
    def rep1(w,g,sm,xb):
        kb = A*3 + B*2
        nb = C*3 - D*w
        pib = sm*kb - nb
        W = pib/4
        for i in range(10):
            for j in range(10):
                xb[i,j] = 1 if W[i,j]>=0 else xb[i,j]
        we = np.sum(w*xb*g)
        return we
    return rep1


g = np.zeros([10,10])
sm = np.ones([10,10])
xb = np.zeros([10,10])
w = 1
r = rep()
rep1 = op(r)
print(rep1(w,g,sm,xb))