code: def test_connect_graphnetwork_nones(self, network_fn):
      batched_graphs_tuple = graph.GraphsTuple(
          n_node=jnp.array([1, 3, 1, 0, 2, 0, 0]),
          n_edge=jnp.array([2, 5, 0, 0, 1, 0, 0]),
          nodes=self._make_nest(jnp.arange(14).reshape(7, 2)),
          edges=self._make_nest(jnp.arange(24).reshape(8, 3)),
          globals=self._make_nest(jnp.arange(14).reshape(7, 2)),
          senders=jnp.array([0, 0, 1, 1, 2, 3, 3, 6]),
          receivers=jnp.array([0, 0, 2, 1, 3, 2, 1, 5]),
      )
  
      for name, graphs_tuple in [
          ("no_globals", batched_graphs_tuple._replace(globals=None)),
          ("empty_globals", batched_graphs_tuple._replace(globals=[])),
          ("no_edges", batched_graphs_tuple._replace(edges=None)),
          ("empty_edges", batched_graphs_tuple._replace(edges=[])),
      ]:
          with self.subTest(name + "_nojit"):
              out = network_fn(graphs_tuple)
              jax.tree_util.tree_map(
                  np.testing.assert_allclose, out, graphs_tuple
              )
          with self.subTest(name + "_jit"):
              out = jax.jit(network_fn)(graphs_tuple)
              jax.tree_util.tree_map(
                  np.testing.assert_allclose, out, 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 4 for calling the ModelsTest obj's underlying member method code with example usage: obj.test_connect_graphnetwork_nones(network_fn) without return types