def index_to_step(index, shape):
  """Compute step for a given nd index if we were enumerating to shape."""
  step = index[0]
  for i, s in enumerate(shape[1:]):
    step = step * s + index[i + 1]
  return step