cupy.fuse#

cupy.fuse(*args, **kwargs)#

用于函数融合的装饰器。

这个装饰器可以比 ElementwiseKernelReductionKernel 更容易地定义逐元素或归约核函数。

由于融合后的核函数会被缓存并重用,建议重用相同的已装饰函数,而不是多次定义并装饰局部函数。

参数:

kernel_name (str) – 融合后的核函数的名称。如果省略,则使用被装饰函数的名称。

示例

>>> @cupy.fuse(kernel_name='squared_diff')
... def squared_diff(x, y):
...     return (x - y) * (x - y)
...
>>> x = cupy.arange(10)
>>> y = cupy.arange(10)[::-1]
>>> squared_diff(x, y)
array([81, 49, 25,  9,  1,  1,  9, 25, 49, 81])