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