ensure_index: raise better error for traced inputs

Author: jakevdp

jax/_src/api_util.py

@@ -37,13 +37,15 @@
 
 def _ensure_index(x: Any) -> Union[int, Tuple[int, ...]]:
   """Ensure x is either an index or a tuple of indices."""
+  x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
   try:
     return operator.index(x)
   except TypeError:
     return tuple(map(operator.index, x))
 
 def _ensure_index_tuple(x: Any) -> Tuple[int, ...]:
   """Convert x to a tuple of indices."""
+  x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
   try:
     return (operator.index(x),)
   except TypeError: