Support output param sharding in dynamo + spmd (#5101)
* Support output param sharding in dynamo + spmd
* Update call to PrepareOutputShardingPropagation and clean up some code
* Run linter
* Fetch coll.indices
* Enable test_dynamo_spmd tests
* Clean up some comments
* Update the PrepareOutputShardingPropagation logic
* Run linter again
* Update replicated_output to default value
* Add virtual device check
* Clean up some code
* Refactor PrepareOutputShardingPropagation into a helper function
* Fix merge conflicts
