device_put_sharded tweaks and docstring detail
device_put_sharded tweaks and docstring detail

I turned some assertions into errors, since they corresponded to
checking whether the user input was valid (rather than internal errors).
I always think of assertions as for internal consistency checks only,
since in principle they can be disabled with `python -O` (though it
seems no one does that).

google/jax,google/jax,google/jax,google/jax