def count_graph_nodes(leaf_nodes: Collection[ValueNode]) -> int:
  """Counts the number of graph nodes.

  Note: these nodes only include the TFT graph nodes, it doesn't count beam
  nodes constructed directly.

  Args:
    leaf_nodes: A list of leaf `ValueNode`s to define the graph.  The graph will
      be the transitive parents of the leaf nodes.

  Returns:
    The count of TFT graph nodes.
  """
  visitor = _CountGraphNodes()
  traverser = Traverser(visitor)
  for value_node in leaf_nodes:
    traverser.visit_value_node(value_node)
  return visitor.num_nodes
Counts the number of graph nodes.

  Note: these nodes only include the TFT graph nodes, it doesn't count beam
  nodes constructed directly.

  Args:
    leaf_nodes: A list of leaf `ValueNode`s to define the graph.  The graph will
      be the transitive parents of the leaf nodes.

  Returns:
    The count of TFT graph nodes.
  
tensorflow_transform/nodes.py