vsel

功能说明

根据condition的数据类型，将condition中的元素与True或1比较，如果表达式成立返回x对应的值，否则返回y对应的值。

• condition的数据类型为bool时，将condition元素的值和True比较，如果表达式成立返回x对应的值，否则返回y对应的值；
• condition的数据类型为uint8时，将condition元素按bit位和1比较，如果表达式成立返回x对应的值，否则返回y对应的值。

以下通过表达式的方式来解释如何比较，其中，i是表示condition中的元素，x'表示x的元素或者x，y'表示y的元素或者y，z表示结果tensor的元素，n（取值0~7）表示condition元素的bit位index，表达式如下：

• condition数据类型为bool：

  z = (i==True)？x'：y'

• condition数据类型为uint8：

  z = (i[n]==1)？x'：y'

  您可以在“te/lang/cce/te_compute/elewise_compute.py”查看接口定义。

约束说明

• x、y元素的数据类型需要一致。
• 当condition的数据类型为uint8时，x、y的shape最后一维要能被8整除。
• 当condition的数据类型为bool时，调用te.lang.cce.build接口进行编译时，传入的config参数中需要设置bool_storage_as_1bit为False。否则会出现某些shape的输出结果不符合预期。

  因为bool_storage_as_1bit的默认值为True，表示按照1 bit进行存储。

  编译配置config如下所示：

  import te.lang.cce as tbe
  
  with tvm.target.cce():
      schedule = tbe.auto_schedule(res)
  config = {"name": kernel_name,
            "tensor_list": [data_x, data_y, res],
            "bool_storage_as_1bit": False}
  tbe.build(schedule, config)

Atlas 200/300/500 推理产品：支持的数据类型有float16

Atlas 训练系列产品：支持的数据类型有float16、float32

函数原型

te.lang.cce.vsel(condition, x, y)

参数说明

• condition：条件tensor，tvm.tensor类型，数据类型为bool、uint8。
• x：可能的返回值，数据类型支持tvm.tensor、scalar。
• y：可能的返回值，数据类型支持tvm.tensor、scalar。

返回值

res_tensor： 结果tensor，tvm.tensor类型

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品