jax

55e77803 - Raise an error when unreduced/reduced is used with the GSPMD partitioner.