【Triton 教程】triton_language.where

Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境，以高效编写自定义 DNN 计算内核，并能够在现代 GPU 硬件上以最大吞吐量运行。

更多 Triton 中文文档可访问 →triton.hyper.ai/

triton.language.where(condition, x, y)

根据condition返回来自x或y元素的张量。
注意：无论condition的值是什么，x和y总是会被求值。
如果希望避免意外的内存操作，请使用triton.load和triton.store中的mask参数。
x和y的形状都会被广播到condition的形状。x和y必须具有相同的数据类型。
参数**：**

• condition（triton.bool 的块）- 当为 True（非零）时，产生 x，否则产生 y。
• x- 在条件为 True 的索引处选择的值。
• y- 在条件为 False 的索引处选择的值。