code: def device_batch(
      graph_generator: data_utils.DataReader,
  ) -> Iterator[jraph.GraphsTuple]:
      """Batches a set of graphs the size of the number of devices."""
      num_devices = jax.local_device_count()
      batch = []
      for idx, graph in enumerate(graph_generator):
          if idx % num_devices == num_devices - 1:
              batch.append(graph)
              yield jax.tree_map(
                  lambda *x: jnp.stack(x, axis=0), *batch
              )
              batch = []
          else:
              batch.append(graph)
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from typing import Iterator
  from absl import app
  from absl import flags
  import haiku
  import jax
  import jax.numpy
  import jraph
  from jraph.ogb_examples import data_utils
  import optax
label_desc: Use image node_id 5 for calling a global function with example usage: device_batch(graph_generator) without return types