Improve error message for indexing into an empty dimension
It seems that JAX (and XLA) are a lot more forgiving for index-out-of-bounds,
except for the case when the dimension has size 0. In this case give a similar
IndexError message as numpy, instead of the current one:
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out
of range, must be within [0, 1), got 1.:
This is a bug in JAX's shape-checking rules; please report it!
Issue: #2671
