0

我试图在 Numba 的 jit 优化代码中使用 numpy，但是当我尝试执行标准 numpy 操作（如 numpy.ones_like）时出现错误，即使 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 我也收到类似的错误

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

4

1 回答 1

1

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

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

于 2020-02-12T18:51:33.590 回答