jax.lax.pmin

jax.lax.pmin(x, axis_name, *, axis_index_groups=None)

在 pmapped 軸 axis_name 上計算 x 的 all-reduce 最小值。

如果 x 是 pytree，則結果等同於將此函式映射到樹狀結構中的每個葉節點。

參數:

• x – 具有名為 axis_name 的映射軸的陣列。

• axis_name – 用於命名 pmapped 軸的可雜湊 Python 物件 (請參閱 jax.pmap() 文件以取得更多詳細資訊)。

• axis_index_groups – 軸索引的選用列表 (例如，對於大小為 4 的軸，[[0, 1], [2, 3]] 將對前兩個和後兩個複本執行 pmin)。群組必須精確地涵蓋所有軸索引一次，並且在 TPU 上，所有群組的大小必須相同。

傳回:

與 x 具有相同形狀的陣列，表示沿軸 axis_name 的 all-reduce 最小值的結果。