def lookup(registered_collection, reg_key):
  """Lookup and return decorated function or class in the collection.

  Lookup decorated function or class in registered_collection, in a
  hierarchical order. For example, when
  reg_key="my_model/my_exp/my_config_0",
  this function will return
  registered_collection["my_model"]["my_exp"]["my_config_0"].

  Args:
    registered_collection: a dictionary. The decorated function or class will be
      retrieved from this collection.
    reg_key: The key for retrieving the registered function or class. If reg_key
      is a string, it can be hierarchical like my_model/my_exp/my_config_0
  Returns:
    The registered function or class.
  Raises:
    LookupError: when reg_key cannot be found.
  """
  if isinstance(reg_key, str):
    hierarchy = reg_key.split("/")
    collection = registered_collection
    for h_idx, entry_name in enumerate(hierarchy):
      if entry_name not in collection:
        raise LookupError(
            "collection path {} at position {} never registered.".format(
                entry_name, h_idx))
      collection = collection[entry_name]
    return collection
  else:
    if reg_key not in registered_collection:
      raise LookupError("registration key {} never registered.".format(reg_key))
    return registered_collection[reg_key]