code: def _batch(graphs, np_):
      """Returns batched graph given a list of graphs and a numpy-like module."""
      # Calculates offsets for sender and receiver arrays, caused by concatenating
      # the nodes arrays.
      offsets = np_.cumsum(
          np_.array([0] + [np_.sum(g.n_node) for g in graphs[:-1]])
      )
  
      def _map_concat(nests):
          concat = lambda *args: np_.concatenate(args)
          return tree.tree_map(concat, *nests)
  
      return gn_graph.GraphsTuple(
          n_node=np_.concatenate([g.n_node for g in graphs]),
          n_edge=np_.concatenate([g.n_edge for g in graphs]),
          nodes=_map_concat([g.nodes for g in graphs]),
          edges=_map_concat([g.edges for g in graphs]),
          globals=_map_concat([g.globals for g in graphs]),
          senders=np_.concatenate(
              [g.senders + o for g, o in zip(graphs, offsets)]
          ),
          receivers=np_.concatenate(
              [g.receivers + o for g, o in zip(graphs, offsets)]
          ),
      )
  
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 14 for calling a global function with example usage: _batch(graphs, np_) and returns: gn_graph, tree