{"id":21481573,"url":"https://github.com/alexoarga/haiku-geometric","last_synced_at":"2025-10-16T11:14:06.203Z","repository":{"id":65149470,"uuid":"581128286","full_name":"alexOarga/haiku-geometric","owner":"alexOarga","description":"A collection of graph neural networks implementations in JAX","archived":false,"fork":false,"pushed_at":"2023-11-28T01:09:02.000Z","size":332,"stargazers_count":31,"open_issues_count":3,"forks_count":1,"subscribers_count":5,"default_branch":"main","last_synced_at":"2024-11-08T12:51:03.315Z","etag":null,"topics":["deep-learning","geometric-deep-learning","graph-neural-networks","jax","machine-learning","python"],"latest_commit_sha":null,"homepage":"https://haiku-geometric.readthedocs.io","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/alexOarga.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-12-22T10:51:03.000Z","updated_at":"2024-11-06T14:05:01.000Z","dependencies_parsed_at":"2023-02-08T06:16:11.763Z","dependency_job_id":null,"html_url":"https://github.com/alexOarga/haiku-geometric","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alexOarga%2Fhaiku-geometric","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alexOarga%2Fhaiku-geometric/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alexOarga%2Fhaiku-geometric/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alexOarga%2Fhaiku-geometric/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/alexOarga","download_url":"https://codeload.github.com/alexOarga/haiku-geometric/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":226041415,"owners_count":17564477,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","geometric-deep-learning","graph-neural-networks","jax","machine-learning","python"],"created_at":"2024-11-23T12:26:51.617Z","updated_at":"2025-10-16T11:14:01.164Z","avatar_url":"https://github.com/alexOarga.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Haiku Geometric\n\n[**Overview**](#overview)\n| [**Installation**](#installation)\n| [**Quickstart**](#quickstart)\n| [**Examples**](#examples)\n| [**Documentation**](https://haiku-geometric.readthedocs.io/en/latest/)\n| [**License**](#license)\n\n[![Documentation Status](https://readthedocs.org/projects/haiku-geometric/badge/?version=latest)](https://haiku-geometric.readthedocs.io/en/latest/?badge=latest)\n[![Python application](https://github.com/alexOarga/haiku-geometric/actions/workflows/python-app.yml/badge.svg)](https://github.com/alexOarga/haiku-geometric/actions/workflows/python-app.yml)\n![pypi](https://img.shields.io/pypi/v/haiku-geometric)\n\n## Overview\n\nHaiku Geometric is a collection of graph neural networks (GNNs) implemented using [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html). It tries to provide **object-oriented** and **easy-to-use** modules for GNNs.\n\nHaiku Geometric is built on top of [Haiku](https://github.com/deepmind/dm-haiku) and [Jraph](https://github.com/deepmind/jraph).\nIt is deeply inspired by [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric). \nIn most cases, Haiku Geometric tries to replicate the API of PyTorch Geometric to allow code sharing between the two.\n\nHaiku Geometric is still under development and I would advise against using it in production.\n\n## Installation\n\nHaiku Geometric can be installed from source:\n\n```bash\npip install git+https://github.com/alexOarga/haiku-geometric.git\n```\n\nAlternatively, you can install Haiku Geometric using pip:\n```bash\npip install haiku-geometric\n```\n\n## Quickstart\n\nFor instance, we can create a simple graph convolutional network (GCN) of 2 layers \nas follows:\n```python\nimport jax\nimport haiku as hk\nfrom haiku_geometric.nn import GCNConv\n\nclass GCN(hk.Module):\n    def __init__(self, hidden_channels, out_channels):\n        super().__init__()\n        self.conv1 = GCNConv(hidden_channels)\n        self.conv2 = GCNConv(hidden_channels)\n        self.linear = hk.Linear(out_channels)\n\n    def __call__(self, nodes,senders, receivers):\n        x = self.conv1(nodes, senders, receivers)\n        x = jax.nn.relu(x)\n        x = self.conv2(x, senders, receivers)\n        x = self.linear(nodes)\n        return x\n\ndef forward(nodes, senders, receivers):\n    gcn = GCN(16, 7)\n    return gcn(nodes, senders, receivers)\n```\n\nThe GNN that we have defined is a Haiku Module. \nTo convert our module in a function that can be used with JAX, we transform\nit using `hk.transform` as described in the \n[Haiku documentation](https://dm-haiku.readthedocs.io/en/latest/).\n\n```python\nmodel = hk.transform(forward)\nmodel = hk.without_apply_rng(model)\nrng = jax.random.PRNGKey(42)\nparams = model.init(rng, nodes=nodes, senders=senders, receivers=receivers)\n```\n\nWe can now run a forward pass on the model:\n```python\noutput = model.apply(params=params, nodes=nodes, senders=senders, receivers=receivers)\n```\n\n## Documentation\n\nThe documentation for Haiku Geometric can be found [here](https://haiku-geometric.readthedocs.io/en/latest/).\n\n## Examples\n\nHaiku Geometric comes with a few examples that showcase the usage of the library.\nThe following examples are available:\n\n|                                                     | Link                                                                                                                                                                                                                                                    |\n|-----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| Quickstart Example                                  | [![Open in Colab](https://img.shields.io/static/v1.svg?logo=google-colab\u0026label=Quickstart\u0026message=Open%20In%20Colab\u0026color=blue)](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/docs/source/notebooks/1_quickstart.ipynb) |\n| Graph Convolution Networks with Karate Club dataset | [![Open in Colab](https://img.shields.io/static/v1.svg?logo=google-colab\u0026label=GCNConv\u0026message=Open%20In%20Colab\u0026color=blue)](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GCNConv_karate_club.ipynb)          |\n| Graph Attention Networks with CORA dataset          | [![Open in Colab](https://img.shields.io/static/v1.svg?logo=google-colab\u0026label=GATConv\u0026message=Open%20In%20Colab\u0026color=blue)](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GATConv_CORA.ipynb)                 |\n| TopKPooling and GraphConv with PROTEINS dataset     | [![Open in Colab](https://img.shields.io/static/v1.svg?logo=google-colab\u0026label=TopKPooling\u0026message=Open%20In%20Colab\u0026color=blue)](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/TopKPooling_GraphConv_PROTEINS.ipynb)                 |\n\n\n## Implemented GNNs modules\n\nCurrently, Haiku Geometric includes the following GNN modules:\n\n| Model                                                                                                                     | Description                                                                                                                                    |\n|---------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|\n| [GCNConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GCNConv)               | Graph convolution layer from the [Semi-Supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907) paper.   |\n| [GATConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GATConv)               | Graph attention layer from the [Graph Attention Networks](https://arxiv.org/abs/1710.10903) paper.                                             |\n| [SAGEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.SAGEConv)             | Graph convolution layer from the [Inductive Representation Learning on Large Graphs](https://arxiv.org/abs/1706.02216) paper.                  |\n| [GINConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINConv)               | Graph isomorphism network layer from the [How Powerful are Graph Neural Networks?](https://arxiv.org/abs/1810.00826) paper.                    |\n| [GINEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINEConv)             | Graph isomorphism network layer from the [Strategies for Pre-training Graph Neural Networks](https://arxiv.org/abs/1905.12265) paper.          |\n| [GraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GraphConv)           | Graph convolution layer from the [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244) paper. |\n| [GeneralConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GeneralConv)       | A general GNN layer adapted from the [Design Space for Graph Neural Networks](https://arxiv.org/abs/2011.08843) paper.                         |\n| [GatedGraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GatedGraphConv) | Graph convolution layer from the [Gated Graph Sequence Neural Networks](https://arxiv.org/abs/1511.05493) paper.                               |\n| [EdgeConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.EdgeConv)             | Edge convolution layer from the [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829) paper.                      |\n| [PNAConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.PNAConv)               | Propagation Network layer from the [Principal Neighbourhood Aggregation for Graph Nets](https://arxiv.org/abs/2004.05718) paper.               |\n| [MetaLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.MetaLayer)           | Meta layer from the [Relational Inductive Biases, Deep Learning, and Graph Networks](https://arxiv.org/abs/1806.01261) paper.                  |\n| [GPSLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GPSLayer)             | Graph layer from the [Recipe for a General, Powerful, Scalable Graph Transformer](https://arxiv.org/abs/2205.12454) paper.                     |\n\n## Implemented positional encodings\n\nThe following positional encodings are currently available:\n\n| Model                                                                                                                     | Description                                                                                                                                     |\n|---------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------|\n| [LaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.LaplacianEncoder)    | Laplacian positional encoding from the [Rethinking Graph Transformers with Spectral Attention](https://arxiv.org/pdf/2106.03893) paper.         |\n| [MagLaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.MagLaplacianEncoder)  | Magnetic Laplacian positional encoding from the [Transformers Meet Directed Graphs](https://arxiv.org/pdf/2302.00049) paper. |\n\n## Issues\n\nIf you encounter any issue, please [open an issue](https://github.com/alexOarga/haiku-geometric/issues/new).\n\n## Running tests\n\nHaiku Geometric can be tested using `pytest` by running the following command:\n\n```bash\npython -m pytest test/\n```\n\n## License\n\n[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Falexoarga%2Fhaiku-geometric","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Falexoarga%2Fhaiku-geometric","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Falexoarga%2Fhaiku-geometric/lists"}