确定 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)

有没有办法在使用 jitclasses 或 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'>))

我认为这是指 python float 和 numba 的 int64 的比较，例如当我做类似 lookup_object.interp(370) 的事情时。

Answer 1

使用type()?

blah = []
if type(blah) is list:
    print "Is a list"

blah = 5
if type(blah) is int:
    print "we have an int"

即：

>>> blah = 5
>>> type(blah)
<type 'int'>
>>>

Answer 2

如果您需要确定和比较 numba jitclass 或 nopython jit 函数中的类型，那么您就不走运了，因为根本不支持 isinstance 并且 type 仅支持少数数字类型和命名元组（请注意，这只是 returns 类型 - 它不适合比较 - 因为 == 没有为里面的 classes 实现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 class 并确定那里的类型，然后相应地转发给 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

Answer 3

从 numba 0.52 开始，支持 np.shape()。所以如果你只想区分 np.ndarray 和标量，下面的工作：

@njit
def test(a):
    if len(np.shape(a)) > 0:
        return 'np.ndarray'
    else:
        return 'not an array'

>>> test(1)
'not an array'
>>> test(np.array([1,2,3]))
'np.ndarray'

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

Determining input argument type for jitclass method

python

numpy

class

numba

编辑