def test_fill_edge_state(self, edge_size):
    """Tests for filling the edge state with a constant content."""
    for g in self.graphs_dicts_in:
      g.pop("edges")
    graphs_tuple = utils_tf.data_dicts_to_graphs_tuple(self.graphs_dicts_in)
    n_edges = np.sum(self.reference_graph.n_edge)
    graphs_tuple = utils_tf.set_zero_edge_features(graphs_tuple, edge_size)
    self.assertAllEqual((n_edges, edge_size),
                        graphs_tuple.edges.get_shape().as_list())
Tests for filling the edge state with a constant content.
graph_nets/tests_tf2/utils_tf_test.py