k3-node 1.0.0__tar.gz

k3_node-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Muhammad Anas Raza
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

k3_node-1.0.0/PKG-INFO

@@ -0,0 +1,284 @@
+Metadata-Version: 2.4
+Name: k3-node
+Version: 1.0.0
+Summary: Multi-Backend Graph Neural Networks on Keras 3
+Author: Muhammad Anas Raza
+License: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: keras>=3.0
+Requires-Dist: scipy
+Requires-Dist: pynndescent
+Requires-Dist: sympy
+Requires-Dist: pandas
+Requires-Dist: huggingface_hub>=0.20.0
+Requires-Dist: onnx>=1.15.0
+Requires-Dist: onnxruntime>=1.17.0
+Requires-Dist: tf2onnx>=1.16.0
+Requires-Dist: onnxscript
+Provides-Extra: examples
+Requires-Dist: scikit-learn; extra == "examples"
+Requires-Dist: rdflib; extra == "examples"
+Requires-Dist: matplotlib; extra == "examples"
+Provides-Extra: test
+Requires-Dist: pytest>=8.0.0; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: torch>=2.0.0; extra == "test"
+Requires-Dist: torch-geometric>=2.5.0; extra == "test"
+Requires-Dist: networkx; extra == "test"
+Requires-Dist: scikit-learn; extra == "test"
+Requires-Dist: tqdm; extra == "test"
+Requires-Dist: fsspec; extra == "test"
+Requires-Dist: requests; extra == "test"
+Requires-Dist: onnx>=1.15.0; extra == "test"
+Requires-Dist: onnxruntime>=1.17.0; extra == "test"
+Requires-Dist: tf2onnx>=1.16.0; extra == "test"
+Requires-Dist: onnxscript; extra == "test"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.5; extra == "docs"
+Requires-Dist: mkdocs-material>=9.5; extra == "docs"
+Requires-Dist: mkdocstrings[python]>=0.25; extra == "docs"
+Requires-Dist: mkdocs-autorefs; extra == "docs"
+Requires-Dist: pymdown-extensions; extra == "docs"
+Requires-Dist: mkdocs-jupyter; extra == "docs"
+Requires-Dist: pygments; extra == "docs"
+Dynamic: license-file
+
+# K3-Node: Multi-Backend Graph Neural Networks
+
+<p align="center">
+  <img src="docs/images/logo.png" alt="K3-Node Logo" width="180"/>
+</p>
+
+<p align="center">
+  <a href="https://anas-rz.github.io/k3-node/"><img src="https://img.shields.io/badge/docs-GitHub%20Pages-blue.svg" alt="Documentation"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_torch.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_torch.yml/badge.svg" alt="PyTorch tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_tensorflow.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_tensorflow.yml/badge.svg" alt="TensorFlow tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_jax.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_jax.yml/badge.svg" alt="JAX tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-MIT-green.svg" alt="License"></a>
+  <a href="https://keras.io/keras_3/"><img src="https://img.shields.io/badge/Keras%203-TensorFlow%20%7C%20PyTorch%20%7C%20JAX-orange.svg" alt="Backends"></a>
+  <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg" alt="Code style: black"></a>
+</p>
+
+---
+
+**K3-Node** is a next-generation graph neural network (GNN) library built natively on **Keras 3**. Write your GNN models once and execute seamlessly across **TensorFlow**, **PyTorch**, and **JAX** with full hardware acceleration (NVIDIA GPUs, Apple Silicon, Google Cloud TPUs).
+
+K3-Node achieves **100% public API parity** with [PyTorch Geometric (PyG)](https://github.com/pyg-team/pytorch_geometric) and incorporates state-of-the-art foundation models and architectures from [Spektral](https://github.com/danielegrattarola/spektral) and [StellarGraph](https://github.com/stellargraph/stellargraph).
+
+📖 **Documentation**: [https://anas-rz.github.io/k3-node/](https://anas-rz.github.io/k3-node/)  
+📋 **Porting Checklist & Parity Status**: [Checklist.md](Checklist.md)
+
+---
+
+## Key Features
+
+- 🔄 **True Multi-Backend Freedom**: Switch between PyTorch, TensorFlow, and JAX with a single environment variable (`KERAS_BACKEND=torch|tensorflow|jax`).
+- 🧠 **Pre-trained Foundation Models**: Out-of-the-box architectures and checkpoint loaders for **GraphMAE2**, **Graphormer** (2D & 3D), **GraphGPS**, **GROVER**, and **Mole-BERT**.
+- ⚡ **65+ Convolution Layers**: Full PyG parity (`GCNConv`, `GATv2Conv`, `TransformerConv`, `GPSConv`, `PNAConv`, `SchNet`, `DimeNetPlusPlus`, `ViSNet`, etc.).
+- 📊 **26 Aggregation Operators**: From elementary aggregations (`sum`, `mean`, `max`, `softmax`, `powermean`) to neural aggregations (`SetTransformer`, `GraphMultisetTransformer`, `Set2Set`, `DeepSets`, `LSTMAggregation`).
+- 🌐 **31 Pooling Operators**: Global readouts (`global_add_pool`, `global_mean_pool`), hierarchical coarsening (`TopKPooling`, `SAGPooling`, `ASAPooling`, `EdgePooling`, `ClusterPooling`), and 3D spatial pooling (`voxel_grid`, `fps`, `knn`, `radius`).
+- 🧱 **Dense & Scalable GNNs**: Dense matrix convolutions (`DenseGCNConv`, `DenseGATConv`), spectral pooling (`DMoNPooling`, `dense_diff_pool`, `dense_mincut_pool`), and linear-complexity graph transformers (`SGFormer`, `LPFormer`, `Polynormer`).
+- 🧭 **Knowledge Graph Embeddings**: Multi-relational link prediction with `TransE`, `RotatE`, `DistMult`, `ComplEx`, and framework-agnostic negative sampling loaders.
+- 📦 **Data, Loaders & Transforms**: Full suite of graph data structures (`Data`, `HeteroData`, `Batch`), mini-batch samplers (`NeighborLoader`, `ClusterLoader`, `GraphSAINTSampler`), and 62+ graph and 3D point cloud transforms.
+- ✅ **Rigorous Verification**: 700+ unit tests on every backend, training tests that check each layer's weights actually learn, compiled-vs-eager and cross-backend consistency tests, and numerical parity tests against PyTorch Geometric and reference checkpoints.
+
+---
+
+## Installation
+
+```bash
+# git should be installed
+pip install git+https://github.com/anas-rz/k3-node/
+
+# with the extra packages the example notebooks use (scikit-learn, rdflib, matplotlib)
+pip install "k3-node[examples] @ git+https://github.com/anas-rz/k3-node"
+```
+
+### Selecting your Backend
+Configure your preferred backend before importing `k3_node`:
+
+```bash
+export KERAS_BACKEND="torch"       # or "tensorflow" or "jax"
+```
+
+Or programmatically in Python:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "torch"  # Must be set before importing k3_node / keras
+import k3_node
+```
+
+---
+
+## Quickstart
+
+### Building a Graph Convolutional Network
+
+```python
+import keras
+from keras import ops
+import k3_node.layers as gnn_layers
+from k3_node.data import Data
+
+class GCN(keras.Model):
+    def __init__(self, in_channels, hidden_channels, out_channels):
+        super().__init__()
+        self.conv1 = gnn_layers.GCNConv(in_channels, hidden_channels)
+        self.conv2 = gnn_layers.GCNConv(hidden_channels, out_channels)
+
+    def call(self, x, edge_index):
+        x = self.conv1(x, edge_index)
+        x = ops.relu(x)
+        x = self.conv2(x, edge_index)
+        return x
+
+# Instantiate model
+model = GCN(in_channels=16, hidden_channels=32, out_channels=7)
+
+# Forward pass on graph data
+x = ops.ones((10, 16))
+edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+
+out = model(x, edge_index)
+print("Output shape:", out.shape)  # (10, 7)
+```
+
+### Training in a Few Lines
+
+The task estimators in `k3_node.tasks` pick the loss, readout and metrics for you:
+
+```python
+from k3_node.datasets import Planetoid
+from k3_node.tasks import NodeClassifier
+
+cora = Planetoid("data/Planetoid", name="Cora")[0]
+
+classifier = NodeClassifier(backbone="gcn", hidden_channels=64, num_layers=2, dropout=0.5)
+classifier.fit(cora, epochs=100, lr=0.01)
+print(classifier.evaluate(cora, mask="test_mask"))
+```
+
+`GraphClassifier`, `GraphRegressor`, `NodeRegressor` and `LinkPredictor` work the same way.
+
+### Example Notebooks
+
+The [`examples/`](examples) folder has 90+ notebooks that follow the architectures of
+[PyG's examples](https://github.com/pyg-team/pytorch_geometric/tree/master/examples), written with
+`keras.Model.fit` and K3-Node's loaders. They cover node, link and graph classification,
+knowledge graphs, molecules (including pre-trained DimeNet, DimeNet++ and SchNet on QM9), point
+clouds, temporal graphs and large-graph mini-batching. Each notebook opens in Colab and runs on
+any backend: change `KERAS_BACKEND` in its first cell. Browse them in the
+[documentation](https://anas-rz.github.io/k3-node/examples/).
+
+---
+
+## Pre-trained Foundation Models
+
+K3-Node provides ready-to-use architectures and automated checkpoint loading for state-of-the-art graph foundation models:
+
+### 1. GraphMAE2 (Self-Supervised Masked Autoencoder)
+```python
+from k3_node.models import GraphMAE2
+from k3_node.models.graphmae2 import load_graphmae2_weights
+
+model = GraphMAE2(
+    in_dim=100,
+    num_hidden=512,
+    out_dim=100,
+    num_layers=4,
+    encoder_type="gat",
+    decoder_type="gat"
+)
+# Load reference pre-trained weights
+load_graphmae2_weights(model, "checkpoints/graphmae2_ogbn_arxiv.pt")
+```
+
+### 2. Graphormer (2D Molecular & 3D Structural Transformer)
+```python
+from k3_node.models import Graphormer, Graphormer3D
+from k3_node.models.graphormer import load_graphormer_weights
+
+# 2D Graphormer (PCQM4Mv2)
+model_2d = Graphormer(num_layers=12, num_heads=32, embed_dim=768)
+load_graphormer_weights(model_2d, "checkpoints/graphormer_pcqm4mv2.pt")
+
+# 3D Graphormer (OC20 Catalyst Adsorption & Molecular Conformations)
+model_3d = Graphormer3D(num_layers=12, num_heads=32, embed_dim=768)
+```
+
+### 3. GraphGPS (Hybrid Local MPNN + Global Transformer)
+```python
+from k3_node.models import GPSModel
+from k3_node.models.gps_model import load_gps_model_weights
+
+model = GPSModel(
+    channels=64,
+    num_layers=5,
+    local_gnn_type="GINE",
+    global_model_type="Transformer"
+)
+load_gps_model_weights(model, "checkpoints/graphgps_zinc.pt")
+```
+
+### 4. GROVER (Self-Supervised Message Passing Transformer)
+```python
+from k3_node.models import GROVER, GROVEREmbedding
+from k3_node.models.grover import load_grover_weights
+
+model = GROVER(hidden_size=128, num_layers=3, num_heads=4)
+load_grover_weights(model, "checkpoints/grover_base.pt")
+```
+
+### 5. Mole-BERT (Masked Chemical Graph Representation)
+```python
+from k3_node.models import MoleBERT
+from k3_node.models.mole_bert import load_mole_bert_weights
+
+model = MoleBERT(num_layer=5, emb_dim=300, drop_ratio=0.5)
+load_mole_bert_weights(model, "checkpoints/Mole-BERT.pth")
+```
+
+---
+
+## What's Included
+
+| Package | Status | Contents |
+|---|---|---|
+| [`k3_node.layers.conv`](https://anas-rz.github.io/k3-node/api/conv/) | ✅ 65/65 | `GCNConv`, `GATConv`, `GATv2Conv`, `SAGEConv`, `GINConv`, `GPSConv`, `TransformerConv`, `PNAConv`, `SchNet`, `DimeNetPlusPlus`, `ViSNet`, etc. |
+| [`k3_node.layers.pool`](https://anas-rz.github.io/k3-node/api/pool/) | ✅ 31/31 | `global_add_pool`, `global_mean_pool`, `TopKPooling`, `SAGPooling`, `ASAPooling`, `EdgePooling`, `ClusterPooling`, `voxel_grid`, `fps`, `graclus`, etc. |
+| [`k3_node.layers.aggr`](https://anas-rz.github.io/k3-node/api/aggr/) | ✅ 26/26 | `SumAggregation`, `MeanAggregation`, `SoftmaxAggregation`, `PowerMeanAggregation`, `MultiAggregation`, `SetTransformerAggregation`, `Set2Set`, etc. |
+| [`k3_node.layers.norm`](https://anas-rz.github.io/k3-node/api/norm/) | ✅ 11/11 | `GraphNorm`, `PairNorm`, `DiffGroupNorm`, `MessageNorm`, `MeanSubtractionNorm`, `BatchNorm`, `LayerNorm`, `HeteroBatchNorm`, etc. |
+| [`k3_node.layers.dense`](https://anas-rz.github.io/k3-node/api/dense/) | ✅ 11/11 | `DenseGCNConv`, `DenseGATConv`, `DenseGINConv`, `DenseSAGEConv`, `DMoNPooling`, `dense_diff_pool`, `dense_mincut_pool`, `Linear`, etc. |
+| [`k3_node.layers.kge`](https://anas-rz.github.io/k3-node/api/kge/) | ✅ 5/5 | `KGEModel`, `TransE`, `RotatE`, `DistMult`, `ComplEx`, `KGTripletLoader`. |
+| [`k3_node.models`](https://anas-rz.github.io/k3-node/api/models/) | ✅ 46/46 | `MLP`, `GAE`, `VGAE`, `DeepGraphInfomax`, `Node2Vec`, `LabelPropagation`, `LINKX`, `LightGCN`, `SGFormer`, `LPFormer`, `Polynormer`, etc. |
+| **Foundation Models** | ✅ 5/5 | `GraphMAE2`, `Graphormer` (2D/3D), `GPSModel`, `GROVER`, `MoleBERT` with pre-trained weight conversion. |
+| [`k3_node.data`](https://anas-rz.github.io/k3-node/api/data/) | ✅ 19/19 | `Data`, `HeteroData`, `Batch`, `TemporalData`, `HypergraphData`, `InMemoryDataset`, `FeatureStore`, `GraphStore`, etc. |
+| [`k3_node.loader`](https://anas-rz.github.io/k3-node/api/loader/) | ✅ 26/26 | `DataLoader`, `NeighborLoader`, `LinkNeighborLoader`, `ClusterLoader`, `GraphSAINTSampler`, `ShaDowKHopSampler`, etc. |
+| [`k3_node.transforms`](https://anas-rz.github.io/k3-node/api/transforms/) | ✅ 62/62 | Topology rewiring, positional encodings (`LapPE`, `RWPE`, `GPSE`), spectral diffusion (`GDC`), and 3D point cloud transforms. |
+
+---
+
+
+## Testing & Verification
+
+Run the comprehensive test suite across backends:
+
+```bash
+# Run all unit tests
+pytest k3_node/
+
+# Run training tests (each layer's weights learn; slower, not run in CI)
+pytest tests_training/
+
+# Run reference parity check against PyTorch implementations
+pytest tests_reference/
+```
+
+---
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

k3_node-1.0.0/README.md

@@ -0,0 +1,237 @@
+# K3-Node: Multi-Backend Graph Neural Networks
+
+<p align="center">
+  <img src="docs/images/logo.png" alt="K3-Node Logo" width="180"/>
+</p>
+
+<p align="center">
+  <a href="https://anas-rz.github.io/k3-node/"><img src="https://img.shields.io/badge/docs-GitHub%20Pages-blue.svg" alt="Documentation"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_torch.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_torch.yml/badge.svg" alt="PyTorch tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_tensorflow.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_tensorflow.yml/badge.svg" alt="TensorFlow tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/actions/workflows/test_jax.yml"><img src="https://github.com/anas-rz/k3-node/actions/workflows/test_jax.yml/badge.svg" alt="JAX tests"></a>
+  <a href="https://github.com/anas-rz/k3-node/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-MIT-green.svg" alt="License"></a>
+  <a href="https://keras.io/keras_3/"><img src="https://img.shields.io/badge/Keras%203-TensorFlow%20%7C%20PyTorch%20%7C%20JAX-orange.svg" alt="Backends"></a>
+  <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg" alt="Code style: black"></a>
+</p>
+
+---
+
+**K3-Node** is a next-generation graph neural network (GNN) library built natively on **Keras 3**. Write your GNN models once and execute seamlessly across **TensorFlow**, **PyTorch**, and **JAX** with full hardware acceleration (NVIDIA GPUs, Apple Silicon, Google Cloud TPUs).
+
+K3-Node achieves **100% public API parity** with [PyTorch Geometric (PyG)](https://github.com/pyg-team/pytorch_geometric) and incorporates state-of-the-art foundation models and architectures from [Spektral](https://github.com/danielegrattarola/spektral) and [StellarGraph](https://github.com/stellargraph/stellargraph).
+
+📖 **Documentation**: [https://anas-rz.github.io/k3-node/](https://anas-rz.github.io/k3-node/)  
+📋 **Porting Checklist & Parity Status**: [Checklist.md](Checklist.md)
+
+---
+
+## Key Features
+
+- 🔄 **True Multi-Backend Freedom**: Switch between PyTorch, TensorFlow, and JAX with a single environment variable (`KERAS_BACKEND=torch|tensorflow|jax`).
+- 🧠 **Pre-trained Foundation Models**: Out-of-the-box architectures and checkpoint loaders for **GraphMAE2**, **Graphormer** (2D & 3D), **GraphGPS**, **GROVER**, and **Mole-BERT**.
+- ⚡ **65+ Convolution Layers**: Full PyG parity (`GCNConv`, `GATv2Conv`, `TransformerConv`, `GPSConv`, `PNAConv`, `SchNet`, `DimeNetPlusPlus`, `ViSNet`, etc.).
+- 📊 **26 Aggregation Operators**: From elementary aggregations (`sum`, `mean`, `max`, `softmax`, `powermean`) to neural aggregations (`SetTransformer`, `GraphMultisetTransformer`, `Set2Set`, `DeepSets`, `LSTMAggregation`).
+- 🌐 **31 Pooling Operators**: Global readouts (`global_add_pool`, `global_mean_pool`), hierarchical coarsening (`TopKPooling`, `SAGPooling`, `ASAPooling`, `EdgePooling`, `ClusterPooling`), and 3D spatial pooling (`voxel_grid`, `fps`, `knn`, `radius`).
+- 🧱 **Dense & Scalable GNNs**: Dense matrix convolutions (`DenseGCNConv`, `DenseGATConv`), spectral pooling (`DMoNPooling`, `dense_diff_pool`, `dense_mincut_pool`), and linear-complexity graph transformers (`SGFormer`, `LPFormer`, `Polynormer`).
+- 🧭 **Knowledge Graph Embeddings**: Multi-relational link prediction with `TransE`, `RotatE`, `DistMult`, `ComplEx`, and framework-agnostic negative sampling loaders.
+- 📦 **Data, Loaders & Transforms**: Full suite of graph data structures (`Data`, `HeteroData`, `Batch`), mini-batch samplers (`NeighborLoader`, `ClusterLoader`, `GraphSAINTSampler`), and 62+ graph and 3D point cloud transforms.
+- ✅ **Rigorous Verification**: 700+ unit tests on every backend, training tests that check each layer's weights actually learn, compiled-vs-eager and cross-backend consistency tests, and numerical parity tests against PyTorch Geometric and reference checkpoints.
+
+---
+
+## Installation
+
+```bash
+# git should be installed
+pip install git+https://github.com/anas-rz/k3-node/
+
+# with the extra packages the example notebooks use (scikit-learn, rdflib, matplotlib)
+pip install "k3-node[examples] @ git+https://github.com/anas-rz/k3-node"
+```
+
+### Selecting your Backend
+Configure your preferred backend before importing `k3_node`:
+
+```bash
+export KERAS_BACKEND="torch"       # or "tensorflow" or "jax"
+```
+
+Or programmatically in Python:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "torch"  # Must be set before importing k3_node / keras
+import k3_node
+```
+
+---
+
+## Quickstart
+
+### Building a Graph Convolutional Network
+
+```python
+import keras
+from keras import ops
+import k3_node.layers as gnn_layers
+from k3_node.data import Data
+
+class GCN(keras.Model):
+    def __init__(self, in_channels, hidden_channels, out_channels):
+        super().__init__()
+        self.conv1 = gnn_layers.GCNConv(in_channels, hidden_channels)
+        self.conv2 = gnn_layers.GCNConv(hidden_channels, out_channels)
+
+    def call(self, x, edge_index):
+        x = self.conv1(x, edge_index)
+        x = ops.relu(x)
+        x = self.conv2(x, edge_index)
+        return x
+
+# Instantiate model
+model = GCN(in_channels=16, hidden_channels=32, out_channels=7)
+
+# Forward pass on graph data
+x = ops.ones((10, 16))
+edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+
+out = model(x, edge_index)
+print("Output shape:", out.shape)  # (10, 7)
+```
+
+### Training in a Few Lines
+
+The task estimators in `k3_node.tasks` pick the loss, readout and metrics for you:
+
+```python
+from k3_node.datasets import Planetoid
+from k3_node.tasks import NodeClassifier
+
+cora = Planetoid("data/Planetoid", name="Cora")[0]
+
+classifier = NodeClassifier(backbone="gcn", hidden_channels=64, num_layers=2, dropout=0.5)
+classifier.fit(cora, epochs=100, lr=0.01)
+print(classifier.evaluate(cora, mask="test_mask"))
+```
+
+`GraphClassifier`, `GraphRegressor`, `NodeRegressor` and `LinkPredictor` work the same way.
+
+### Example Notebooks
+
+The [`examples/`](examples) folder has 90+ notebooks that follow the architectures of
+[PyG's examples](https://github.com/pyg-team/pytorch_geometric/tree/master/examples), written with
+`keras.Model.fit` and K3-Node's loaders. They cover node, link and graph classification,
+knowledge graphs, molecules (including pre-trained DimeNet, DimeNet++ and SchNet on QM9), point
+clouds, temporal graphs and large-graph mini-batching. Each notebook opens in Colab and runs on
+any backend: change `KERAS_BACKEND` in its first cell. Browse them in the
+[documentation](https://anas-rz.github.io/k3-node/examples/).
+
+---
+
+## Pre-trained Foundation Models
+
+K3-Node provides ready-to-use architectures and automated checkpoint loading for state-of-the-art graph foundation models:
+
+### 1. GraphMAE2 (Self-Supervised Masked Autoencoder)
+```python
+from k3_node.models import GraphMAE2
+from k3_node.models.graphmae2 import load_graphmae2_weights
+
+model = GraphMAE2(
+    in_dim=100,
+    num_hidden=512,
+    out_dim=100,
+    num_layers=4,
+    encoder_type="gat",
+    decoder_type="gat"
+)
+# Load reference pre-trained weights
+load_graphmae2_weights(model, "checkpoints/graphmae2_ogbn_arxiv.pt")
+```
+
+### 2. Graphormer (2D Molecular & 3D Structural Transformer)
+```python
+from k3_node.models import Graphormer, Graphormer3D
+from k3_node.models.graphormer import load_graphormer_weights
+
+# 2D Graphormer (PCQM4Mv2)
+model_2d = Graphormer(num_layers=12, num_heads=32, embed_dim=768)
+load_graphormer_weights(model_2d, "checkpoints/graphormer_pcqm4mv2.pt")
+
+# 3D Graphormer (OC20 Catalyst Adsorption & Molecular Conformations)
+model_3d = Graphormer3D(num_layers=12, num_heads=32, embed_dim=768)
+```
+
+### 3. GraphGPS (Hybrid Local MPNN + Global Transformer)
+```python
+from k3_node.models import GPSModel
+from k3_node.models.gps_model import load_gps_model_weights
+
+model = GPSModel(
+    channels=64,
+    num_layers=5,
+    local_gnn_type="GINE",
+    global_model_type="Transformer"
+)
+load_gps_model_weights(model, "checkpoints/graphgps_zinc.pt")
+```
+
+### 4. GROVER (Self-Supervised Message Passing Transformer)
+```python
+from k3_node.models import GROVER, GROVEREmbedding
+from k3_node.models.grover import load_grover_weights
+
+model = GROVER(hidden_size=128, num_layers=3, num_heads=4)
+load_grover_weights(model, "checkpoints/grover_base.pt")
+```
+
+### 5. Mole-BERT (Masked Chemical Graph Representation)
+```python
+from k3_node.models import MoleBERT
+from k3_node.models.mole_bert import load_mole_bert_weights
+
+model = MoleBERT(num_layer=5, emb_dim=300, drop_ratio=0.5)
+load_mole_bert_weights(model, "checkpoints/Mole-BERT.pth")
+```
+
+---
+
+## What's Included
+
+| Package | Status | Contents |
+|---|---|---|
+| [`k3_node.layers.conv`](https://anas-rz.github.io/k3-node/api/conv/) | ✅ 65/65 | `GCNConv`, `GATConv`, `GATv2Conv`, `SAGEConv`, `GINConv`, `GPSConv`, `TransformerConv`, `PNAConv`, `SchNet`, `DimeNetPlusPlus`, `ViSNet`, etc. |
+| [`k3_node.layers.pool`](https://anas-rz.github.io/k3-node/api/pool/) | ✅ 31/31 | `global_add_pool`, `global_mean_pool`, `TopKPooling`, `SAGPooling`, `ASAPooling`, `EdgePooling`, `ClusterPooling`, `voxel_grid`, `fps`, `graclus`, etc. |
+| [`k3_node.layers.aggr`](https://anas-rz.github.io/k3-node/api/aggr/) | ✅ 26/26 | `SumAggregation`, `MeanAggregation`, `SoftmaxAggregation`, `PowerMeanAggregation`, `MultiAggregation`, `SetTransformerAggregation`, `Set2Set`, etc. |
+| [`k3_node.layers.norm`](https://anas-rz.github.io/k3-node/api/norm/) | ✅ 11/11 | `GraphNorm`, `PairNorm`, `DiffGroupNorm`, `MessageNorm`, `MeanSubtractionNorm`, `BatchNorm`, `LayerNorm`, `HeteroBatchNorm`, etc. |
+| [`k3_node.layers.dense`](https://anas-rz.github.io/k3-node/api/dense/) | ✅ 11/11 | `DenseGCNConv`, `DenseGATConv`, `DenseGINConv`, `DenseSAGEConv`, `DMoNPooling`, `dense_diff_pool`, `dense_mincut_pool`, `Linear`, etc. |
+| [`k3_node.layers.kge`](https://anas-rz.github.io/k3-node/api/kge/) | ✅ 5/5 | `KGEModel`, `TransE`, `RotatE`, `DistMult`, `ComplEx`, `KGTripletLoader`. |
+| [`k3_node.models`](https://anas-rz.github.io/k3-node/api/models/) | ✅ 46/46 | `MLP`, `GAE`, `VGAE`, `DeepGraphInfomax`, `Node2Vec`, `LabelPropagation`, `LINKX`, `LightGCN`, `SGFormer`, `LPFormer`, `Polynormer`, etc. |
+| **Foundation Models** | ✅ 5/5 | `GraphMAE2`, `Graphormer` (2D/3D), `GPSModel`, `GROVER`, `MoleBERT` with pre-trained weight conversion. |
+| [`k3_node.data`](https://anas-rz.github.io/k3-node/api/data/) | ✅ 19/19 | `Data`, `HeteroData`, `Batch`, `TemporalData`, `HypergraphData`, `InMemoryDataset`, `FeatureStore`, `GraphStore`, etc. |
+| [`k3_node.loader`](https://anas-rz.github.io/k3-node/api/loader/) | ✅ 26/26 | `DataLoader`, `NeighborLoader`, `LinkNeighborLoader`, `ClusterLoader`, `GraphSAINTSampler`, `ShaDowKHopSampler`, etc. |
+| [`k3_node.transforms`](https://anas-rz.github.io/k3-node/api/transforms/) | ✅ 62/62 | Topology rewiring, positional encodings (`LapPE`, `RWPE`, `GPSE`), spectral diffusion (`GDC`), and 3D point cloud transforms. |
+
+---
+
+
+## Testing & Verification
+
+Run the comprehensive test suite across backends:
+
+```bash
+# Run all unit tests
+pytest k3_node/
+
+# Run training tests (each layer's weights learn; slower, not run in CI)
+pytest tests_training/
+
+# Run reference parity check against PyTorch implementations
+pytest tests_reference/
+```
+
+---
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

k3_node-1.0.0/k3_node/__init__.py

@@ -0,0 +1,122 @@
+"""
+`k3_node` is a library for building multibackend graph neural networks. 
+Built upon Keras 3.0 the models can be trained using TensorFlow, PyTorch, 
+or JAX.
+
+To install the package, run:
+    
+```bash
+git clone https://github.com/anas-rz/k3-node.git # bash
+```
+        
+```python
+# in your code
+import sys
+sys.path.append('k3-node')
+
+import os
+os.environ['KERAS_BACKEND'] = 'tensorflow' # or 'torch' or 'jax'
+
+from k3_node import ...
+```
+"""
+__version__ = "1.0.0"
+
+import warnings
+
+# Index tensors are requested as int64; without `jax_enable_x64`, JAX stores them as int32
+# (which is all graphs of this size need) and would warn on every conversion.
+warnings.filterwarnings("ignore", message="Explicitly requested dtype int64", category=UserWarning)
+
+from k3_node import data
+from k3_node.data import Data, Batch
+from k3_node import datasets
+from k3_node import io
+from k3_node import layers
+from k3_node import loader
+from k3_node import transforms
+from k3_node import models
+from k3_node import applications
+
+from k3_node import metrics
+from k3_node import tasks
+from k3_node.tasks import (
+    NodeClassifier,
+    NodeRegressor,
+    GraphClassifier,
+    GraphRegressor,
+    LinkPredictor,
+)
+
+from k3_node import etl
+from k3_node.etl import (
+    TableToGraph,
+    TabularToGraph,
+    table_to_graph,
+    RelationalToGraph,
+    relational_to_graph,
+)
+
+from k3_node import hub
+from k3_node.hub import (
+    from_pretrained,
+    push_to_hub,
+    save_pretrained,
+    load_dataset_from_hub,
+    push_dataset_to_hub,
+)
+
+from k3_node import export
+from k3_node.export import (
+    export_onnx,
+    export_tflite,
+    export_tensorrt,
+    generate_triton_config,
+    ONNXModel,
+    TFLiteModel,
+)
+
+from k3_node import rag
+from k3_node.layers import kge
+
+__all__ = [
+    "data",
+    "Data",
+    "Batch",
+    "datasets",
+    "io",
+    "layers",
+    "loader",
+    "transforms",
+    "models",
+    "applications",
+    "metrics",
+    "tasks",
+    "NodeClassifier",
+    "NodeRegressor",
+    "GraphClassifier",
+    "GraphRegressor",
+    "LinkPredictor",
+    "etl",
+    "TableToGraph",
+    "TabularToGraph",
+    "table_to_graph",
+    "RelationalToGraph",
+    "relational_to_graph",
+    "hub",
+    "from_pretrained",
+    "push_to_hub",
+    "save_pretrained",
+    "load_dataset_from_hub",
+    "push_dataset_to_hub",
+    "export",
+    "export_onnx",
+    "export_tflite",
+    "export_tensorrt",
+    "generate_triton_config",
+    "ONNXModel",
+    "TFLiteModel",
+    "rag",
+    "kge",
+]
+

k3_node-1.0.0/k3_node/applications/__init__.py

@@ -0,0 +1,17 @@
+"""Domain-specific applications built on K3-Node.
+
+Submodules:
+- `chemistry`: Molecular graphs, quantum property prediction, SMILES processing, Uni-Mol, AttentiveFP, SchNet, DimeNet.
+- `bio`: Macromolecular structures, protein-protein interactions, docking, UniMolDocking.
+- `materials`: Crystal graph neural networks, periodic boundaries, CHGNet, M3GNet, MEGNet, TensorNet, SO3Net, QET.
+"""
+
+from k3_node.applications import bio
+from k3_node.applications import chemistry
+from k3_node.applications import materials
+
+__all__ = [
+    "bio",
+    "chemistry",
+    "materials",
+]