TIK 1.5自定义算子开发指南(Beta)-refine_shape_axes:调用示例

时间:2023-11-01 16:13:42

调用示例

from tbe.common.utils import shape_utilshape, axes= shape_util.refine_shape_axes((2, 3, 4, 5, 6),(1,-3)) 
  1. 首先对axes(1,-3)进行负值优化,由于shape(2, 3, 4, 5, 6)共有5维,所以-3为倒数第三个维度,即2,则优化后的reduce列表axes=(1, 2)。
  2. 对shape(2, 3, 4, 5, 6)进行fuse操作:

    对连续的reduce轴1、轴2进行fuse操作,对连续的非reduce轴3、轴4进行fuse操作,操作完成后shape=(2, 12,30),axes=(1, 2) fuse成了轴1。

    所以返回值的shape与axes如下:

    shape=(2, 12, 30)

    axes=(1)

support.huaweicloud.com/TIKopdevgd_beta/atlasoputilsapi_07_0018.html