code: def setUp(self):
      super().setUp()
      self.rngs = {
          "params": jax.random.key(0),
          "dropout": jax.random.key(1),
      }
      n_node = jnp.arange(3, 11)
      n_edge = jnp.arange(4, 12)
      total_n_node = jnp.sum(n_node)
      total_n_edge = jnp.sum(n_edge)
      n_graph = n_node.shape[0]
      feature_dim = 10
      self.graphs = jraph.GraphsTuple(
          n_node=n_node,
          n_edge=n_edge,
          senders=jnp.zeros(total_n_edge, dtype=jnp.int32),
          receivers=jnp.ones(total_n_edge, dtype=jnp.int32),
          nodes=jnp.ones((total_n_node, feature_dim)),
          edges=jnp.zeros((total_n_edge, feature_dim)),
          globals=jnp.zeros((n_graph, feature_dim)),
      )
  
file_path: flax/examples/ogbg_molpcba/models_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import jraph
  import models
label_desc: Use image node_id 1 for calling the ModelsTest obj's underlying member method code with example usage: obj.setUp() without return types