Helper to determine the shape of reorder output.

def _reorder_shape(input_shape, output=None):  # pylint: disable=invalid-name
  """Helper to determine the shape of reorder output."""
  if output is None:
    return input_shape
  return base.nested_map(output, lambda i: input_shape[i])