code: def test_fully_connected_graph(
      self, n_node, n_graph, nodes, globals_
  ):
      node_feat = (
          np.random.rand(n_node * n_graph, 32) if nodes else None
      )
      global_feat = np.random.rand(n_graph, 32) if globals_ else None
      with self.subTest("nojit"):
          result = utils.get_fully_connected_graph(
              n_node, n_graph, node_feat, global_feat
          )
          if nodes:
              self.assertLen(result.nodes, n_node * n_graph)
          if globals_:
              self.assertLen(result.globals, n_graph)
          self.assertLen(result.senders, n_node**2 * n_graph)
          self.assertLen(result.receivers, n_node**2 * n_graph)
          np.testing.assert_allclose(
              result.n_node, jnp.array([n_node] * n_graph)
          )
      with self.subTest("jit"):
          result = jax.jit(
              utils.get_fully_connected_graph, static_argnums=[0, 1]
          )(n_node, n_graph, node_feat, global_feat)
          if nodes:
              self.assertLen(result.nodes, n_node * n_graph)
          if globals_:
              self.assertLen(result.globals, n_graph)
          self.assertLen(result.senders, n_node**2 * n_graph)
          self.assertLen(result.receivers, n_node**2 * n_graph)
          np.testing.assert_allclose(
              result.n_node, jnp.array([n_node] * n_graph)
          )
  
      with self.subTest("senders_receiver_indices"):
          if n_node > 0:
              # [0, 1, ..., n_node - 1]
              node_indices = jnp.arange(n_node)
              # [0, 1,..., n_node - 1] + [0, 1,..., n_node - 1] + ... n_node times
              # [0,..., 0, 1,..., 1,..., n_node - 1,..., n_node - 1] each n_node times
              expected_senders = np.concatenate(
                  [node_indices] * n_node, axis=0
              )
              expected_receivers = np.stack(
                  [node_indices] * n_node, axis=-1
              ).reshape([-1])
          else:
              expected_senders = np.array([], dtype=np.int32)
              expected_receivers = np.array([], dtype=np.int32)
  
          # Check sender and receivers on each graph in the batch.
          for result_graph in utils.unbatch(result):
              np.testing.assert_allclose(
                  result_graph.senders, expected_senders
              )
              np.testing.assert_allclose(
                  result_graph.receivers, expected_receivers
              )
  
imports: import functools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax.lib import xla_bridge
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 29 for calling the GraphTest obj's underlying member method code with example usage: obj.test_fully_connected_graph(n_node, n_graph, nodes, globals_) without return types