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])