triton.language.max

triton.language.max(input, axis=None, return_indices=False, return_indices_tie_break_left=True, keep_dims=False)

返回 input 张量沿指定 axis 的所有元素的最大值

参数:
  • input (Tensor) – 输入值

  • axis (int) – 应执行归约的维度。如果为 None,则归约所有维度

  • keep_dims (bool) – 如果为 true,保留长度为 1 的归约维度

  • return_indices (bool) – 如果为 true,返回与最大值对应的索引

  • return_indices_tie_break_left (bool) – 如果为 true,在存在并列(即,多个元素具有相同的最大值)的情况下,返回非 NaN 值的最左侧索引

此函数也可以作为成员函数在 tensor 上调用,形如 x.max(...) 而非 max(x, ...)