code: def test_connect_gnns(self, network_fn):
      batched_graphs_tuple = graph.GraphsTuple(
          n_node=jnp.array([1, 3, 1, 0, 2, 0, 0]),
          n_edge=jnp.array([1, 7, 1, 0, 3, 0, 0]),
          nodes=jnp.arange(14).reshape(7, 2),
          edges=jnp.arange(36).reshape(12, 3),
          globals=jnp.arange(14).reshape(7, 2),
          senders=jnp.array([0, 1, 2, 3, 4, 5, 6, 1, 2, 3, 3, 6]),
          receivers=jnp.array([0, 1, 2, 3, 4, 5, 6, 2, 3, 2, 1, 5]),
      )
      with self.subTest("nojit"):
          out, expected_out = network_fn(batched_graphs_tuple)
          jax.tree_util.tree_map(
              np.testing.assert_allclose, out, expected_out
          )
      with self.subTest("jit"):
          out, expected_out = jax.jit(network_fn)(batched_graphs_tuple)
          jax.tree_util.tree_map(
              np.testing.assert_allclose, out, expected_out
          )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import models
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 5 for calling the ModelsTest obj's underlying member method code with example usage: obj.test_connect_gnns(network_fn) without return types