code: def setUp(self):
      super().setUp()
      os.environ[
          "XLA_FLAGS"
      ] = "--xla_force_host_platform_device_count=3"
      xla_bridge.get_backend.cache_clear()
  
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 the ShardedGraphnetTest obj's underlying member method code with example usage: obj.setUp() without return types