jax

ed33d102 - Add ppermute as an allowed multi-host collective. (#1981)

Add ppermute as an allowed multi-host collective. (#1981) I manually tested that this works as of 0417e1e. The indices used in ppermute correspond to those returned by `axis_index`.