code: def test_dynamically_batch(self, use_globals, batch_kwargs):
      def graph_iterator():
          graphs = [
              _make_dynamic_batch_graph(add_globals=use_globals)
              for x in range(4)
          ]
          return iter(graphs + utils.unbatch_np(graphs[-1]))
  
      batched_dataset = utils.dynamically_batch(
          graph_iterator(), **batch_kwargs
      )
      batched_graphs = []
      while True:
          try:
              batched_graphs.append(next(batched_dataset))
          except StopIteration:
              break
  
      self.assertLen(batched_graphs, 5)
      for batch_graphs in batched_graphs:
          batch_nodes = jax.tree_util.tree_flatten(batch_graphs.nodes)[
              0
          ]
          for nodes in batch_nodes:
              self.assertEqual(nodes.shape[0], batch_kwargs["n_node"])
          batch_edges = jax.tree_util.tree_flatten(batch_graphs.edges)[
              0
          ]
          for edges in batch_edges:
              self.assertEqual(edges.shape[0], batch_kwargs["n_edge"])
          self.assertLen(batch_graphs.n_node, batch_kwargs["n_graph"])
          self.assertEqual(
              utils.get_number_of_padding_with_graphs_nodes(
                  batch_graphs
              ),
              batch_kwargs["n_node"] - sum(_DB_NUM_NODES),
          )
          self.assertEqual(
              utils.get_number_of_padding_with_graphs_edges(
                  batch_graphs
              ),
              batch_kwargs["n_edge"] - sum(_DB_NUM_EDGES),
          )
  
imports: import functools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax.lib import xla_bridge
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 2 for calling the DynamicBatchTest obj's underlying member method code with example usage: obj.test_dynamically_batch(use_globals, batch_kwargs) and returns: iter