code: def test_connect_graphnetwork(self, network_fn):
      _, batched_graphs_tuple = self._get_list_and_batched_graph()
      with self.subTest("nojit"):
          out = network_fn(batched_graphs_tuple)
          jax.tree_util.tree_map(
              np.testing.assert_allclose, out, batched_graphs_tuple
          )
      with self.subTest("jit"):
          out = jax.jit(network_fn)(batched_graphs_tuple)
          jax.tree_util.tree_map(
              np.testing.assert_allclose, out, batched_graphs_tuple
          )
  
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 3 for calling the ModelsTest obj's underlying member method code with example usage: obj.test_connect_graphnetwork(network_fn) without return types