transpose

功能说明

将输入tensor按照axes的值转置，详细功能介绍可参见调用示例。

函数原型

transpose(tensor, axes)

参数说明

• tensor：输入tensor，tvm.tensor类型。

  支持的数据类型：in8, uint8, float16, int16, uint16, float32, int32, uint32, int64, uint64。

• axes：对输入tensor的转置索引，需要为int类型的list或者tuple。axes的长度为d，d是输入tensor的维度。axes中每一个元素的取值范围为[0, d-1]，且每一个元素的取值不能重复。

返回值

res_tensor：转置之后的tensor，tvm.tensor类型。

约束说明

此接口暂不支持与其他TBE DSL计算接口混合使用。

支持的型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

Atlas 推理系列产品

Atlas 200/500 A2推理产品

Atlas A2训练系列产品/Atlas 800I A2推理产品

调用示例

输入Tensor x为:

x = [[1,2,3],
     [4,5,6]]

输入axes为：

axes = [1,0]

调用transpose接口对x行列转置，调用示例如下：

from tbe import tvm
from tbe import dsl
shape = (2,3) 
input_dtype = "int32"
x = tvm.placeholder(shape, name="x", dtype=input_dtype)  
axes = [1,0]
res = dsl.transpose(x, axes) 
# res.shape = (3,2)

以上代码的功能为将输入Tensor x的轴0与轴1进行转置，输出Tensor res为：

res = [[1,4],
       [2,5],
       [3,6]]