jax

5b8775dc - [sharding_in_types] Add sharding rule for reduce sum which is just drop the specs for the axis we are reducing over