def _extract_sparse_output_annotations(
    graph: tf.compat.v1.Graph) -> List[Tuple[tf.Tensor, List[tf.Tensor]]]:
  """Extracts sparse output annotations from the graph."""
  tensors = graph.get_collection(_ANNOTATED_SPARSE_SHAPE_TENSORS)
  shapes = graph.get_collection(_ANNOTATED_SPARSE_SHAPES)
  assert len(tensors) == len(shapes), f'{tensors} != {shapes}'
  return list(zip(tensors, shapes))
Extracts sparse output annotations from the graph.
tensorflow_transform/schema_inference.py