code: def test_sharded_same_as_non_sharded(self, n_edge):
      in_tuple = _get_graphs_from_n_edge(n_edge)
      devices = 3
      sharded_tuple = sharded_graphnet.graphs_tuple_to_broadcasted_sharded_graphs_tuple(
          in_tuple, devices
      )
      update_fn = jraph.concatenated_args(lambda x: x)
      sharded_gn = sharded_graphnet.ShardedEdgesGraphNetwork(
          update_fn, update_fn, update_fn, num_shards=devices
      )
      gn = jraph.GraphNetwork(update_fn, update_fn, update_fn)
      sharded_out = jax.pmap(sharded_gn, axis_name="i")(sharded_tuple)
      expected_out = gn(in_tuple)
      reduced_out = sharded_graphnet.broadcasted_sharded_graphs_tuple_to_graphs_tuple(
          sharded_out
      )
      jax.tree_util.tree_map(
          functools.partial(
              np.testing.assert_allclose, atol=1e-5, rtol=1e-5
          ),
          expected_out,
          reduced_out,
      )
  
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 3 for calling the ShardedGraphnetTest obj's underlying member method code with example usage: obj.test_sharded_same_as_non_sharded(n_edge) without return types