triton.language.softmax

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

计算 x 的元素级 Softmax。

参数:

x (Block) – 输入值

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