triton.language.softmax

triton.language.softmax(x, dim=None, keep_dims=False, ieee_rounding=False)

计算 x 的逐元素 softmax。

参数:

x (Block) – 输入值

此函数也可以作为 tensor 的成员函数调用，形式为 x.softmax(...)，而不是 softmax(x, ...)。