code: def estimate_padding_budget_for_batch_size(
      dataset: tf.data.Dataset,
      batch_size: int,
      num_estimation_graphs: int,
  ) -> GraphsTupleSize:
      """Estimates the padding budget for a dataset of unbatched GraphsTuples.
  
      Args:
        dataset: A dataset of unbatched GraphsTuples.
        batch_size: The intended batch size. Note that no batching is performed by
          this function.
        num_estimation_graphs: How many graphs to take from the dataset to estimate
          the distribution of number of nodes and edges per graph.
  
      Returns:
        padding_budget: The padding budget for batching and padding the graphs
        in this dataset to the given batch size.
      """
  
      def next_multiple_of_64(val: float):
          """Returns the next multiple of 64 after val."""
          return 64 * (1 + int(val // 64))
  
      if batch_size <= 1:
          raise ValueError(
              "Batch size must be > 1 to account for padding graphs."
          )
  
      total_num_nodes = 0
      total_num_edges = 0
      for graph in dataset.take(
          num_estimation_graphs
      ).as_numpy_iterator():
          graph_size = get_graphs_tuple_size(graph)
          if graph_size.n_graph != 1:
              raise ValueError("Dataset contains batched GraphTuples.")
  
          total_num_nodes += graph_size.n_node
          total_num_edges += graph_size.n_edge
  
      num_nodes_per_graph_estimate = (
          total_num_nodes / num_estimation_graphs
      )
      num_edges_per_graph_estimate = (
          total_num_edges / num_estimation_graphs
      )
  
      padding_budget = GraphsTupleSize(
          n_node=next_multiple_of_64(
              num_nodes_per_graph_estimate * batch_size
          ),
          n_edge=next_multiple_of_64(
              num_edges_per_graph_estimate * batch_size
          ),
          n_graph=batch_size,
      )
      return padding_budget
  
file_path: flax/examples/ogbg_molpcba/input_pipeline.py
imports: import functools
  from typing import Dict, NamedTuple
  import jraph
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: estimate_padding_budget_for_batch_size(dataset, batch_size, num_estimation_graphs) and returns: padding_budget, unknown