docstr: Tests padding of graph.
docstr_tok: ["Tests","padding","of","graph","."]
code: def test_pad_with_graphs(self):
      """Tests padding of graph."""
      _, graphs_tuple = _get_list_and_batched_graph()
      padded_graphs_tuple = utils.pad_with_graphs(
          graphs_tuple, 10, 12, 9
      )
      expected_padded_graph = graph.GraphsTuple(
          n_node=jnp.concatenate(
              [graphs_tuple.n_node, jnp.array([3, 0])]
          ),
          n_edge=jnp.concatenate(
              [graphs_tuple.n_edge, jnp.array([4, 0])]
          ),
          nodes=tree.tree_map(
              lambda f: jnp.concatenate(
                  [f, jnp.zeros((3, 2), dtype=f.dtype)]
              ),
              graphs_tuple.nodes,
          ),
          edges=tree.tree_map(
              lambda f: jnp.concatenate(
                  [f, jnp.zeros((4, 3), dtype=f.dtype)]
              ),
              graphs_tuple.edges,
          ),
          globals=tree.tree_map(
              lambda f: jnp.concatenate(
                  [f, jnp.zeros((2, 2), dtype=f.dtype)]
              ),
              graphs_tuple.globals,
          ),
          senders=jnp.concatenate(
              [graphs_tuple.senders, jnp.array([7, 7, 7, 7])]
          ),
          receivers=jnp.concatenate(
              [graphs_tuple.receivers, jnp.array([7, 7, 7, 7])]
          ),
      )
      jax.tree_util.tree_map(
          np.testing.assert_allclose,
          padded_graphs_tuple,
          expected_padded_graph,
      )
  
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 6 for calling the GraphTest obj's underlying member method code with example usage: obj.test_pad_with_graphs() without return types