Schema visualization tool

Schemas can be complex and need visualization. A simple script could convert the schema to a networkx graph and generate an image with the graph visualization.

This can be a starting point:

import networkx as nx
import matplotlib.pyplot as plt


graph_schema = tfgnn.read_schema("schema.pbtxt")
graph = nx.DiGraph()

nodes = [
    name for name in graph_schema.node_sets.keys()
]
edges = [
    [e.source, e.target] for e in graph_schema.edge_sets.values()
]
edge_labels = {
    (e.source, e.target): name for name, e in graph_schema.edge_sets.items()
}

graph.add_nodes_from(nodes)
graph.add_edges_from(edges)

pos = nx.spring_layout(graph)

nx.draw(graph,
        pos=pos,
        node_size=1000,
        with_labels = True)

nx.draw_networkx_edge_labels(
    graph,
    pos=pos,
    edge_labels=edge_labels
)

plt.savefig("schema.png")

The above can generate a schema visualizations like this:

The issue names a proposed Python script that reads schema.pbtxt and generates schema.png, but it does not identify a repository file, entry point, or test. Start by locating how schemas are currently read and where a visualization tool would belong. Done means producing an image with labeled schema nodes and edges from the input schema.