def _populate_test_data(self, max_size):
    """Populates the class fields with data used for the tests.

    This creates a batch of graphs with number of nodes from 0 to `num`,
    number of edges ranging from 1 to `num`, plus an empty graph with no nodes
    and no edges (so that the total number of graphs is 1 + (num ** (num + 1)).

    The nodes states, edges states and global states of the graphs are
    created to have different types and shapes.

    Those graphs are stored both as dictionaries (in `self.graphs_dicts_in`,
    without `n_node` and `n_edge` information, and in `self.graphs_dicts_out`
    with these two fields filled), and a corresponding numpy
    `graphs.GraphsTuple` is stored in `self.reference_graph`.

    Args:
      max_size: The maximum number of nodes and edges (inclusive).
    """
    filt = lambda x: (x[0] > 0) or (x[1] == 0)
    n_node, n_edge = zip(*list(
        filter(filt, itertools.product(
            range(max_size + 1), range(max_size + 1)))))

    graphs_dicts = []
    nodes = []
    edges = []
    receivers = []
    senders = []
    globals_ = []

    def _make_default_state(shape, dtype):
      return np.arange(np.prod(shape)).reshape(shape).astype(dtype)

    for i, (n_node_, n_edge_) in enumerate(zip(n_node, n_edge)):
      n = _make_default_state([n_node_, 7, 11], "f4") + i * 100.
      e = _make_default_state([n_edge_, 13, 14], np.float64) + i * 100. + 1000.
      r = _make_default_state([n_edge_], np.int32) % n_node[i]
      s = (_make_default_state([n_edge_], np.int32) + 1) % n_node[i]
      g = _make_default_state([5, 3], "f4") - i * 100. - 1000.

      nodes.append(n)
      edges.append(e)
      receivers.append(r)
      senders.append(s)
      globals_.append(g)
      graphs_dict = dict(nodes=n, edges=e, receivers=r, senders=s, globals=g)
      graphs_dicts.append(graphs_dict)

    # Graphs dicts without n_node / n_edge (to be used as inputs).
    self.graphs_dicts_in = graphs_dicts
    # Graphs dicts with n_node / n_node (to be checked against outputs).
    self.graphs_dicts_out = []
    for dict_ in self.graphs_dicts_in:
      completed_dict = dict_.copy()
      completed_dict["n_node"] = completed_dict["nodes"].shape[0]
      completed_dict["n_edge"] = completed_dict["edges"].shape[0]
      self.graphs_dicts_out.append(completed_dict)

    # pylint: disable=protected-access
    offset = utils_np._compute_stacked_offsets(n_node, n_edge)
    # pylint: enable=protected-access
    self.reference_graph = graphs.GraphsTuple(**dict(
        nodes=np.concatenate(nodes, axis=0),
        edges=np.concatenate(edges, axis=0),
        receivers=np.concatenate(receivers, axis=0) + offset,
        senders=np.concatenate(senders, axis=0) + offset,
        globals=np.stack(globals_),
        n_node=np.array(n_node),
        n_edge=np.array(n_edge)))
Populates the class fields with data used for the tests.

    This creates a batch of graphs with number of nodes from 0 to `num`,
    number of edges ranging from 1 to `num`, plus an empty graph with no nodes
    and no edges (so that the total number of graphs is 1 + (num ** (num + 1)).

    The nodes states, edges states and global states of the graphs are
    created to have different types and shapes.

    Those graphs are stored both as dictionaries (in `self.graphs_dicts_in`,
    without `n_node` and `n_edge` information, and in `self.graphs_dicts_out`
    with these two fields filled), and a corresponding numpy
    `graphs.GraphsTuple` is stored in `self.reference_graph`.

    Args:
      max_size: The maximum number of nodes and edges (inclusive).
    
graph_nets/tests_tf2/test_utils.py