def test_correct_signature(
      self,
      dynamic_num_nodes,
      dynamic_num_edges,
      dynamic_num_graphs,
      batched,
      replace_globals_with_constant,
      nested_features_type):
    """Tests that the correct spec is created when using different options."""

    if batched:
      input_data_dicts = [self.graphs_dicts[1], self.graphs_dicts[2]]
    else:
      input_data_dicts = [self.graphs_dicts[1]]

    graph = utils_np.data_dicts_to_graphs_tuple(input_data_dicts)
    num_graphs = len(input_data_dicts)
    num_edges = sum(graph.n_edge).item()
    num_nodes = sum(graph.n_node).item()

    # Manually setting edges and globals fields to give some variety in
    # testing situations.
    # Making edges have rank 1 to .
    graph = graph.replace(edges=np.zeros(num_edges))

    # Make a constant field.
    if replace_globals_with_constant:
      graph = graph.replace(globals=np.array(0.0, dtype=np.float32))

    if nested_features_type is not None:
      graph = self._add_nested_features(graph, nested_features_type)

    spec_signature = utils_tf.specs_from_graphs_tuple(
        graph, dynamic_num_graphs, dynamic_num_nodes, dynamic_num_edges)

    # Captures if nodes/edges will be dynamic either due to dynamic nodes/edges
    # or dynamic graphs.
    dynamic_nodes_or_graphs = dynamic_num_nodes or dynamic_num_graphs
    dynamic_edges_or_graphs = dynamic_num_edges or dynamic_num_graphs

    num_edges = None if dynamic_edges_or_graphs else num_edges
    num_nodes = None if dynamic_nodes_or_graphs else num_nodes
    num_graphs = None if dynamic_num_graphs else num_graphs

    if replace_globals_with_constant:
      expected_globals_shape = []
    else:
      expected_globals_shape = [num_graphs,] + test_utils.GLOBALS_DIMS

    expected_answer = graphs.GraphsTuple(
        nodes=tf.TensorSpec(
            shape=[num_nodes,] + test_utils.NODES_DIMS,
            dtype=tf.float32),
        edges=tf.TensorSpec(
            shape=[num_edges],  # Edges were manually replaced to have dim 1.
            dtype=tf.float64),
        n_node=tf.TensorSpec(
            shape=[num_graphs],
            dtype=tf.int32),
        n_edge=tf.TensorSpec(
            shape=[num_graphs],
            dtype=tf.int32),
        globals=tf.TensorSpec(
            shape=expected_globals_shape,
            dtype=tf.float32),
        receivers=tf.TensorSpec(
            shape=[num_edges],
            dtype=tf.int32),
        senders=tf.TensorSpec(
            shape=[num_edges],
            dtype=tf.int32),
        )

    if nested_features_type is not None:
      expected_answer = self._add_nested_features(
          expected_answer, nested_features_type)

    with self.subTest(name="Correct Type."):
      self.assertIsInstance(spec_signature, graphs.GraphsTuple)

    with self.subTest(name="Correct Structure."):
      tree.assert_same_structure(spec_signature, expected_answer)

    with self.subTest(name="Correct Signature."):
      def assert_equal(actual, expected):
        self.assertEqual(actual, expected)
        return True
      tree.map_structure(assert_equal, spec_signature, expected_answer)
Tests that the correct spec is created when using different options.
graph_nets/tests/utils_tf_test.py