Variational Graph Autoencoder

JAX implementation of Variational Graph Auto-Encoders (Kipf and Welling, 2016).

Dependencies

• JAX
• Jraph
• Haiku
• Optax
• NetworkX
• scikit-learn

Usage

Run either GAE (Graph Autoencoder) or VGAE (Variational GAE):

python3 train.py --is_vgae=True

Todos

• Make a runnable notebook with loss plots, etc.
• Try on different datasets
• Compare ROC-AUC with results in Kipf-Welling paper
• Add documentation comments for encoder, decoder, etc.

---

Attribution: A good chunk of the code, especially those related to dataset preprocessing, is attributed to this amazing Colab tutorial by Lisa Wang and Nikola Jovanović.