In some cases 0d tensors have no sharding. This PR provides a few minor fixes to account for such cases.