96399c3ec6
* Fix sharding when no device_map is passed * style * add tests * align * add docstring * format --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>