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

我认为这是指当我做类似

float

的事情时，python

int64

和 numba

lookup_object.interp(370)

的比较。

Answer 1

如果您需要确定和比较 numba

jitclass

或 nopython

jit

函数内的类型，那么您就不走运了，因为根本不支持

isinstance

并且

type

仅支持少数数字类型，并且命名元组（请注意，这只是返回类型 - 它不适合比较 - 因为

==

没有为 numba 函数内的类实现）。

从 Numba 0.35 开始，唯一支持的内置函数是（来源：numba 文档）：

支持以下内置功能：

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.njit

ted 函数：

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 2

从 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'

Answer 3

也许有点晚了，但你可以尝试使用

objmode

：

@njit
def isarray(obj):
    with objmode(isarray="boolean"):
        isarray = isinstance(obj, np.ndarray)
    return isarray

然后使用

isarray(x0)

代替

isinstance(x0, np.ndarray)

。

Answer 4

这是根据参数类型替换所需函数实现的示例。

在 numba 0.57.1 上测试

import numba as nb
import numba.experimental as nbexp
import numba.extending as nbex
from numba import types as nbt

@nbexp.jitclass([ ('_x', nbt.float32),
                  ('_y', nbt.float32), ])
class Vec2:
    def __init__(self, x : float, y : float):
        self._x = x
        self._y = y

    @property
    def x(self) -> float: return self._x
    @property
    def y(self) -> float: return self._y

    def __mul__(self, other): return Vec2(0,0) # overloaded

    
# Overload implementations
def Vec2__mul__Vec2(self, other):   return Vec2(self._x*other._x, self._y*other._y)
def Vec2__mul__number(self, other): return Vec2(self._x*float(other), self._y*float(other))

# Overloaders
@nbex.overload_method(nbt.misc.ClassInstanceType, "__mul__")
def over_Vec2__mul__(self, other):
    if self is Vec2.class_type.instance_type:
        if other is Vec2.class_type.instance_type: return Vec2__mul__Vec2
        if other in nbt.number_domain: return Vec2__mul__number

# Tests

@nb.njit(nogil=True)
def run_test1():
    return Vec2(1,1) * 2

@nb.njit(nogil=True)
def run_test2():
    return Vec2(1,1) * Vec2(3,3)

print( run_test1().x ) # outputs 2.0
print( run_test2().x ) # outputs 3.0

Answer 5

使用

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'>
>>>