def get_structure(dataset_or_iterator):
  """Returns the type specification of an element of a `Dataset` or `Iterator`.

  Args:
    dataset_or_iterator: A `tf.data.Dataset` or `tf.data.Iterator`.

  Returns:
    A nested structure of `tf.TypeSpec` objects matching the structure of an
    element of `dataset_or_iterator` and spacifying the type of individal
    components.

  Raises:
    TypeError: If `dataset_or_iterator` is not a `Dataset` or `Iterator` object.
  """
  try:
    return dataset_or_iterator.element_spec  # pylint: disable=protected-access
  except AttributeError:
    raise TypeError("`dataset_or_iterator` must be a Dataset or Iterator "
                    "object, but got %s." % type(dataset_or_iterator))