jax

acbca16b - Error out if sorted dimension is sharded because that almost always entails an AG. It's better to reshard at the user level instead.