【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 的索引处选择的值。