使用内部带有 for 循环的 jitted 函数

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

这似乎与评论中提到的@max9111 有关 - xb 数组是在 jitted 函数之外创建的，因此无法在内部修改。我通过将 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))

使用内部带有 for 循环的 jitted 函数

Using a jitted function with a for loop inside

python

jit

loops

for-loop

numba