code: def _get_graphs_from_n_edge(n_edge):
      """Get a graphs tuple from n_edge."""
      graphs = []
      for el in n_edge:
          graphs.append(
              jraph.GraphsTuple(
                  nodes=np.random.uniform(size=(128, 2)),
                  edges=np.random.uniform(size=(el, 2)),
                  senders=np.random.choice(128, el),
                  receivers=np.random.choice(128, el),
                  n_edge=np.array([el]),
                  n_node=np.array([128]),
                  globals=np.array([[el]]),
              )
          )
      graphs = utils.batch_np(graphs)
      return graphs
  
file_path: jraph/jraph/experimental/sharded_graphnet_test.py
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.tree_util
  import jraph
  from jraph._src import utils
  from jraph.experimental import sharded_graphnet
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _get_graphs_from_n_edge(n_edge) and returns: graphs