count += 1
    for function in tf_graph_def.library.function:
      for node_def in function.node_def:
        if node_def.op == 'XlaCallModule':
          count += 1
    return count

  def _get_function_aliases(
      self, output_saved_model_path: str, tags: List[str]
  ) -> dict[str, str]:
    """Gets the function aliases in the output saved model."""
    loader = loader_impl.SavedModelLoader(output_saved_model_path)