Package: haiku-geometric
Haiku GeometricOverview|Installation|Quickstart|Examples|Documentation|LicenseOverviewHaiku Geometric is a collection of graph neural networks (GNNs) implemented usingJAX. It tries to provideobject-orientedandeasy-to-usemodules for GNNs.Haiku Geometric is built on top ofHaikuandJraph.
It is deeply inspired byPyTorch 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.InstallationHaiku Geometric can be installed from source:pipinstallgit+https://github.com/alexOarga/haiku-geometric.gitAlternatively, you can install Haiku Geometric using pip:pipinstallhaiku-geometricQuickstartFor instance, we can create a simple graph convolutional network (GCN) of 2 layers
as follows:importjaximporthaikuashkfromhaiku_geometric.nnimportGCNConvclassGCN(hk.Module):def__init__(self,hidden_channels,out_channels):super().__init__()self.conv1=GCNConv(hidden_channels)self.conv2=GCNConv(hidden_channels)self.linear=hk.Linear(out_channels)def__call__(self,nodes,senders,receivers):x=self.conv1(nodes,senders,receivers)x=jax.nn.relu(x)x=self.conv2(x,senders,receivers)x=self.linear(nodes)returnxdefforward(nodes,senders,receivers):gcn=GCN(16,7)returngcn(nodes,senders,receivers)The GNN that we have defined is a Haiku Module.
To convert our module in a function that can be used with JAX, we transform
it usinghk.transformas described in theHaiku documentation.model=hk.transform(forward)model=hk.without_apply_rng(model)rng=jax.random.PRNGKey(42)params=model.init(rng,nodes=nodes,senders=senders,receivers=receivers)We can now run a forward pass on the model:output=model.apply(params=params,nodes=nodes,senders=senders,receivers=receivers)DocumentationThe documentation for Haiku Geometric can be foundhere.ExamplesHaiku Geometric comes with a few examples that showcase the usage of the library.
The following examples are available:LinkQuickstart ExampleGraph Convolution Networks with Karate Club datasetGraph Attention Networks with CORA datasetTopKPooling and GraphConv with PROTEINS datasetImplemented GNNs modulesCurrently, Haiku Geometric includes the following GNN modules:ModelDescriptionGCNConvGraph convolution layer from theSemi-Supervised Classification with Graph Convolutional Networkspaper.GATConvGraph attention layer from theGraph Attention Networkspaper.SAGEConvGraph convolution layer from theInductive Representation Learning on Large Graphspaper.GINConvGraph isomorphism network layer from theHow Powerful are Graph Neural Networks?paper.GINEConvGraph isomorphism network layer from theStrategies for Pre-training Graph Neural Networkspaper.GraphConvGraph convolution layer from theWeisfeiler and Leman Go Neural: Higher-order Graph Neural Networkspaper.GeneralConvA general GNN layer adapted from theDesign Space for Graph Neural Networkspaper.GatedGraphConvGraph convolution layer from theGated Graph Sequence Neural Networkspaper.EdgeConvEdge convolution layer from theDynamic Graph CNN for Learning on Point Cloudspaper.PNAConvPropagation Network layer from thePrincipal Neighbourhood Aggregation for Graph Netspaper.MetaLayerMeta layer from theRelational Inductive Biases, Deep Learning, and Graph Networkspaper.GPSLayerGraph layer from theRecipe for a General, Powerful, Scalable Graph Transformerpaper.Implemented positional encodingsThe following positional encodings are currently available:ModelDescriptionLaplacianEncoderLaplacian positional encoding from theRethinking Graph Transformers with Spectral Attentionpaper.MagLaplacianEncoderMagnetic Laplacian positional encoding from theTransformers Meet Directed Graphspaper.IssuesIf you encounter any issue, pleaseopen an issue.Running testsHaiku Geometric can be tested usingpytestby running the following command:python-mpytesttest/License