Revert "Refactor shard_args ShardedDeviceArray slow path. (#2562)" (#2577)
Revert "Refactor shard_args ShardedDeviceArray slow path. (#2562)" (#2577)

This reverts commit 2e87674aff0609d58eed07091d8e9fad8ee71d6c.

This refactoring was broken, but we apparently have no test coverage!
mypy/linting would have caught it though.
google/jax,tensorflow/probability,google/jax,google/jax,google/jax,tensorflow/probability