code: def get_graphs_tuples(n_edge, sharded_n_edge, device_graph_idx):
      sharded_n_edge = np.array(sharded_n_edge)
      device_graph_idx = np.array(device_graph_idx)
      devices = len(sharded_n_edge)
      graphs = _get_graphs_from_n_edge(n_edge)
      sharded_senders = np.reshape(graphs.senders, [devices, -1])
      sharded_receivers = np.reshape(graphs.receivers, [devices, -1])
      sharded_edges = np.reshape(
          graphs.edges, [devices, -1, graphs.edges.shape[-1]]
      )
      # Broadcast replicated features to have a devices leading axis.
      broadcast = lambda x: np.broadcast_to(
          x[None, :], [devices] + list(x.shape)
      )
  
      sharded_graphs = sharded_graphnet.ShardedEdgesGraphsTuple(
          device_senders=sharded_senders,
          device_receivers=sharded_receivers,
          device_edges=sharded_edges,
          device_n_edge=sharded_n_edge,
          nodes=broadcast(graphs.nodes),
          senders=broadcast(graphs.senders),
          receivers=broadcast(graphs.receivers),
          device_graph_idx=device_graph_idx,
          globals=broadcast(graphs.globals),
          n_node=broadcast(graphs.n_node),
          n_edge=broadcast(graphs.n_edge),
      )
      return graphs, sharded_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 2 for calling a global function with example usage: get_graphs_tuples(n_edge, sharded_n_edge, device_graph_idx) and returns: graphs,  sharded_graphs