jax.lax.pmin#
- jax.lax.pmin(x, axis_name, *, axis_index_groups=None)[source]#
Compute an all-reduce min on
x
over the pmapped axisaxis_name
.If
x
is a pytree then the result is equivalent to mapping this function to each leaf in the tree.- Parameters:
x – array(s) with a mapped axis named
axis_name
.axis_name – hashable Python object used to name a pmapped axis (see the
jax.pmap()
documentation for more details).axis_index_groups – optional list of lists containing axis indices (e.g. for an axis of size 4, [[0, 1], [2, 3]] would perform pmins over the first two and last two replicas). Groups must cover all axis indices exactly once, and on TPUs all groups must be the same size.
- Returns:
Array(s) with the same shape as
x
representing the result of an all-reduce min along the axisaxis_name
.