使用Numba的异常-使用不受支持的NumPy函数'numpy.ones_like'或不受支持的使用函数

Question

我正在尝试从Numba的jit优化代码中使用numpy，但是当我尝试执行numpy.ones_like之类的标准numpy操作时遇到错误，即使numba文档中提到支持该操作。

文档链接：Numba 0.46。

代码：

import numba
from numba import jit
import pandas as pd
import numpy as np

print(numba.__version__)

@jit(nopython=True)
def calc_method(a,b):
    a1 = np.float64(a)
    b1 = np.float64(b)
    abc = (a1, np.ones_like(b1))
    abc_ht = np.hstack(abc)
    return abc_ht

def calculate(cudf_df: cudf, size_of_row: int):       
    return cudf_df.apply_chunks(calc_method, incols=['a', 'b'], outcols=dict(), chunks=size_of_row)

df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7, 8], 'b': [11, 12, 13, 14, 15, 16, 17, 18]})
cudf_df = cudf.DataFrame.from_pandas(df)
a, b = calculate(cudf_df, 4)

错误：

TypingError                               Traceback (most recent call last)
<ipython-input-38-ad56fb75bc4a> in <module>
----> 1 a, b = calculate(cudf_df, 4)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<numba.cuda.compiler.DeviceFunctionTemplate object at 0x7fa78521b550>) with argument(s) of type(s): (array(int64, 1d, A), array(int64, 1d, A))
 * parameterized
In definition 0:
    TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Use of unsupported NumPy function 'numpy.ones_like' or unsupported use of the function.

File "<ipython-input-37-97f7d707ba81>", line 9:
def calc_method(a,b):
    <source elided>
    b1 = np.float64(b)
    abc = (a1, np.ones_like(b1))
    ^

有人可以告诉我在上面的示例中我在做什么错吗？预先感谢。

我也对np.hstack遇到类似的错误

注意：这是重现此问题的简化示例。

Answer 1

您不能使用任何从JIT内核中分配内存的numpy方法。通常，您需要提前分配输出，然后在内核中设置这些输出的值。

您可以在此处看到使用apply_chunks的示例：https://gist.github.com/beckernick/acbfb9e8ac4f0657789930a0dfb57d17#file-udf_apply_chunks_basic_example-ipynb