Data Structure coding for GDAS NAS

how to code the following graph structure ?
which kind of data structure is suitable in this case for both forward inference and backward propagation ?
Would Tutorial — NetworkX 3.2.1 documentation or PyG Documentation — pytorch_geometric documentation or https://www.dgl.ai/ be suitable ?

https://gist.github.com/promach/b6f526c56e20f029d68e6f9041c3f5c0#file-gdas-py-L106-L107