gather

功能说明

获取输入Tensor的特定维度中指定索引的切片数据。

调用此接口,则代表需要在params的第0维中,分别取索引为2和索引为0的数,结果如下所示:

gather_tensor = [[7,8,9], [1,2,3]]

函数原型

gather(params, indices, axis=None, batch_dims=0)

参数说明

返回值

wrapped_tensor:执行gather之后的Tensor,tvm.Tensor类型

约束说明

若配置了batch_dims,即batch_dims≠0,则params与indices的batch轴的维度大小要保持一致。例如,若batch_dims=1,表示要在第一个轴上做批处理,则params与indices的轴0的维度大小要相同。

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例