code: def test_get_sharded_graphs_tuple(
      self, n_edge, sharded_n_edge, device_graph_idx
  ):
      in_tuple, expect_tuple = get_graphs_tuples(
          n_edge, sharded_n_edge, device_graph_idx
      )
      out_tuple = sharded_graphnet.graphs_tuple_to_broadcasted_sharded_graphs_tuple(
          in_tuple, num_shards=len(expect_tuple.nodes)
      )
      tree.tree_map(
          np.testing.assert_almost_equal, out_tuple, expect_tuple
      )
  
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 the ShardedGraphnetTest obj's underlying member method code with example usage: obj.test_get_sharded_graphs_tuple(n_edge, sharded_n_edge, device_graph_idx) without return types