def METHOD_NAME(  # pylint: disable=invalid-name
    key_maybe_dict: Union[str, Dict[str, str]],
    features: Dict[str,
                   np.ndarray]) -> Tuple[Optional[List[str]], Optional[Union[
                       np.ndarray, Dict[str, np.ndarray]]]]:
  """Returns keys and values in dict given key (or dict of keys)."""
  if isinstance(key_maybe_dict, dict):
    values = {}
    keys = set()
    for output_name, key in key_maybe_dict.items():
      if key in features:
        values[output_name] = features[key]
        if key not in keys:
          keys.add(key)
    return (list(keys), values)
  elif key_maybe_dict in features:
    return ([key_maybe_dict], features[key_maybe_dict])
  else:
    return ([], None)