Haiku Geometric Documentation

Graph Neural Networks in JAX

Haiku Geometric is a collection of graph neural network (GNN) implementations in JAX. It tries to provide object-oriented and easy-to-use modules for GNNs.

Haiku Geometric is built on top of Haiku and Jraph. It is deeply inspired by PyTorch Geometric. In most cases, Haiku Geometric tries to replicate the API of PyTorch Geometric to allow code sharing between the two.

Haiku Geometric is still under development and I would advise against using it in production.

Installation

Haiku Geomtric can be installed via Pypi:

pip install haiku-geometric

Alternatively you can install it from source:

pip install git+https://github.com/alexOarga/haiku-geometric

Contents

• Quickstart
• Creating your own dataset
• Batch support
• Example Notebooks

API Reference

• haiku_geometric.nn
• haiku_geometric.datasets
• haiku_geometric.models
• haiku_geometric.posenc
• haiku_geometric.transforms
• haiku_geometric.utils

Indices and tables

• Index

• Module Index

• Search Page