code: def broadcasted_sharded_graphs_tuple_to_graphs_tuple(
      sharded_graphs_tuple,
  ):
      """Converts a broadcasted ShardedGraphsTuple to a GraphsTuple."""
      # We index the first element of replicated arrays, since they have been
      # repeated. For edges, we reshape to recover all of the edge features.
      unbroadcast = lambda y: tree.tree_map(lambda x: x[0], y)
      unshard = lambda x: jnp.reshape(
          x, (x.shape[0] * x.shape[1],) + x.shape[2:]
      )
      # TODO(jonathangodwin): check senders and receivers are consistent.
      return jraph.GraphsTuple(
          nodes=unbroadcast(sharded_graphs_tuple.nodes),
          edges=tree.tree_map(
              unshard, sharded_graphs_tuple.device_edges
          ),
          n_node=sharded_graphs_tuple.n_node[0],
          n_edge=sharded_graphs_tuple.n_edge[0],
          globals=unbroadcast(sharded_graphs_tuple.globals),
          senders=sharded_graphs_tuple.senders[0],
          receivers=sharded_graphs_tuple.receivers[0],
      )
  
file_path: jraph/jraph/experimental/sharded_graphnet.py
imports: import functools
  from typing import Callable, List, NamedTuple, Optional
  import jax
  import jax.numpy
  import jax.tree_util
  import jraph
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: broadcasted_sharded_graphs_tuple_to_graphs_tuple(sharded_graphs_tuple) and returns: jraph