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

Question

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

文档链接：Numba 0.46。

编辑：如果我直接调用方法"calc_method"，则工作正常，从apply_chunks内使用时失败。所以可能不是 Numba 本身的问题，而是如何使用cudf.apply_chunks。

法典：

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