code: def test_estimate_padding_budget_valid(self, valid_batch_size):
      budget = input_pipeline.estimate_padding_budget_for_batch_size(
          self.datasets["train"],
          valid_batch_size,
          num_estimation_graphs=1,
      )
      self.assertEqual(budget.n_graph, valid_batch_size)
  
file_path: flax/examples/ogbg_molpcba/input_pipeline_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import input_pipeline
  import jraph
  import tensorflow
label_desc: Use image node_id 2 for calling the InputPipelineTest obj's underlying member method code with example usage: obj.test_estimate_padding_budget_valid(valid_batch_size) without return types