python - 确定 jitclass 方法的输入参数类型

Question

我正在开发一个 jitclass，其中的方法之一可以接受 int、float 或 numpy.ndarray 的输入参数>。我需要能够确定参数是数组还是其他两种类型中的任何一种。我尝试使用 isinstance，如下面的 interp 方法所示:

spec = [('x', float64[:]),
        ('y', float64[:])]


@jitclass(spec)
class Lookup:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def interp(self, x0):
        if isinstance(x0, (float, int)):
            result = self._interpolate(x0)
        elif isinstance(x0, np.ndarray):
            result = np.zeros(x0.size)
            for i in range(x0.size):
                result[i] = self._interpolate(x0[i])
        else:
            raise TypeError("`interp` method can only accept types of float, int, or ndarray.")
        return result

    def _interpolate(self, x0):
        x = self.x
        y = self.y
        if x0 < x[0]:
            return y[0]
        elif x0 > x[-1]:
            return y[-1]
        else:
            for i in range(len(x) - 1):
                if x[i] <= x0 <= x[i + 1]:
                    x1, x2 = x[i], x[i + 1]
                    y1, y2 = y[i], y[i + 1]

                    return y1 + (y2 - y1) / (x2 - x1) * (x0 - x1)

但我收到以下错误:

numba.errors.TypingError: Failed at nopython (nopython frontend)
Failed at nopython (nopython frontend)
Untyped global name 'isinstance': cannot determine Numba type of <class 'builtin_function_or_method'>
File "Lookups.py", line 17
[1] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'interp') for instance.jitclass.Lookup#2167664ca28<x:array(float64, 1d, A),y:array(float64, 1d, A)>)
[2] During: typing of call at <string> (3)

在使用 jitclass 或 nopython 模式时，有办法确定输入参数是否属于某种类型吗？

编辑

我应该之前提到过这一点，但使用内置的 type 似乎也不起作用。例如，如果我将 interp 方法替换为:

def interp(self, x0):
        if type(x0) == float or type(x0) == int:
            result = self._interpolate(x0)
        elif type(x0) == np.ndarray:
            result = np.zeros(x0.size)
            for i in range(x0.size):
                result[i] = self._interpolate(x0[i])
        else:
            raise TypeError("`interp` method can only accept types of float, int, or ndarray.")
        return result

我收到以下错误:

numba.errors.TypingError: Failed at nopython (nopython frontend)
Failed at nopython (nopython frontend)
Invalid usage of == with parameters (class(int64), Function(<class 'float'>))

我认为这是指当我执行诸如 lookup_object.interp(370) 之类的操作时，python float 和 numba 的 int64 的比较示例。

Answer 1

如果您需要确定和比较 numba jitclass 或 nopython jit 函数内的类型，那么您就不走运了，因为 isinstance 不是根本不支持，并且 type 仅支持少数数字类型和命名元组(请注意，这仅返回类型 - 它不适合比较 - 因为 == 不是” t 为 numba 函数内的类实现)。

从 Numba 0.35 开始，唯一支持的内置函数是(来源:numba documentation):

The following built-in functions are supported:

abs()
bool
complex
divmod()
enumerate()
float
int: only the one-argument form
iter(): only the one-argument form
len()
min()
max()
next(): only the one-argument form
print(): only numbers and strings; no file or sep argument
range: semantics are similar to those of Python 3 even in Python 2: a range object is returned instead of an array of values.
round()
sorted(): the key argument is not supported
type(): only the one-argument form, and only on some types (e.g. numbers and named tuples)
zip()

我的建议:使用普通的 Python 类并确定其中的类型，然后相应地转发到 numba.njitted 函数:

import numba as nb
import numpy as np

@nb.njit
def _interpolate_one(x, y, x0):
    if x0 < x[0]:
        return y[0]
    elif x0 > x[-1]:
        return y[-1]
    else:
        for i in range(len(x) - 1):
            if x[i] <= x0 <= x[i + 1]:
                x1, x2 = x[i], x[i + 1]
                y1, y2 = y[i], y[i + 1]

                return y1 + (y2 - y1) / (x2 - x1) * (x0 - x1)

@nb.njit
def _interpolate_many(x, y, x0):
    result = np.zeros(x0.size, dtype=np.float_)
    for i in range(x0.size):
        result[i] = _interpolate_one(x, y, x0[i])
    return result

class Lookup:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def interp(self, x0):
        if isinstance(x0, (float, int)):
            result = _interpolate_one(self.x, self.y, x0)
        elif isinstance(x0, np.ndarray):
            result = _interpolate_many(self.x, self.y, x0)
        else:
            raise TypeError("`interp` method can only accept types of float, int, or ndarray.")
        return result