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 2.5 documentation or PyTorch Geometric Documentation — pytorch_geometric 1.7.0 documentation or https://www.dgl.ai/ be suitable ?