code: def unbatch(
      graph: gn_graph.GraphsTuple,
  ) -> List[gn_graph.GraphsTuple]:
      """Returns a list of graphs given a batched graph.
  
      This function does not support jax.jit, because the shape of the output
      is data-dependent!
  
      Args:
        graph: the batched graph, which will be unbatched into a list of graphs.
      """
      return _unbatch(graph, np_=jnp)
  
imports: import functools
  from typing import Any, Callable, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Union
  import jax
  from jax import lax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  import numpy
label_desc: Use image node_id 15 for calling a global function with example usage: unbatch(graph) and returns: _unbatch