Modify clipping implementation to avoid jnp.moveaxis, which causes undesirable all-to-all's in distributed environments.

Modify clipping implementation to avoid jnp.moveaxis, which causes undesirable all-to-all's in distributed environments. #4629