code: def test_graph_net(
      self,
      latent_size: int,
      output_globals_size: int,
      use_edge_model: bool,
  ):
      # Input definition.
      graphs = self.graphs
      num_nodes = jnp.sum(graphs.n_node)
      num_edges = jnp.sum(graphs.n_edge)
      num_graphs = graphs.n_node.shape[0]
  
      # Model definition.
      net = models.GraphNet(
          latent_size=latent_size,
          num_mlp_layers=2,
          message_passing_steps=2,
          output_globals_size=output_globals_size,
          use_edge_model=use_edge_model,
      )
      output, _ = net.init_with_output(self.rngs, graphs)
  
      # Output should be graph with the same topology, but a
      # different number of features.
      self.assertIsInstance(output, jraph.GraphsTuple)
      self.assertSequenceAlmostEqual(output.n_node, graphs.n_node)
      self.assertSequenceAlmostEqual(output.n_edge, graphs.n_edge)
      self.assertSequenceAlmostEqual(output.senders, graphs.senders)
      self.assertSequenceAlmostEqual(output.receivers, graphs.receivers)
      self.assertEqual(output.nodes.shape, (num_nodes, latent_size))
      self.assertEqual(output.edges.shape, (num_edges, latent_size))
      self.assertEqual(
          output.globals.shape, (num_graphs, output_globals_size)
      )
  
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 3 for calling the ModelsTest obj's underlying member method code with example usage: obj.test_graph_net(latent_size, output_globals_size, use_edge_model) without return types