Aggregation Layers

The k3_node.layers.aggr module provides 26 basic, statistical, scaled, and neural aggregation operators for neighborhood messaging and graph readout.

Aggregation Base Class

`k3_node.layers.aggr.Aggregation`

Bases: Layer

An abstract base class for implementing custom aggregations.

`reduce(x, index=None, ptr=None, dim_size=None, dim=-2, reduce='sum')`

Reduces features along groups specified by index or ptr.

`reset_parameters()`

Resets all learnable parameters of the module.

Basic Aggregations

SumAggregation

`k3_node.layers.aggr.SumAggregation`

Bases: Aggregation

An aggregation operator that sums up features across a set of elements.

Example

import numpy as np
from k3_node.layers import SumAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = SumAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

MeanAggregation

`k3_node.layers.aggr.MeanAggregation`

Bases: Aggregation

An aggregation operator that averages features across a set of elements.

Example

import numpy as np
from k3_node.layers import MeanAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MeanAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

MaxAggregation

`k3_node.layers.aggr.MaxAggregation`

Bases: Aggregation

An aggregation operator that takes the feature-wise maximum across a set of elements.

Example

import numpy as np
from k3_node.layers import MaxAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MaxAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

MinAggregation

`k3_node.layers.aggr.MinAggregation`

Bases: Aggregation

An aggregation operator that takes the feature-wise minimum across a set of elements.

Example

import numpy as np
from k3_node.layers import MinAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MinAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

MulAggregation

`k3_node.layers.aggr.MulAggregation`

Bases: Aggregation

An aggregation operator that multiplies features across a set of elements.

Example

import numpy as np
from k3_node.layers import MulAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MulAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

Statistical & Scaled Aggregations

SoftmaxAggregation

`k3_node.layers.aggr.SoftmaxAggregation`

Bases: Aggregation

The softmax aggregation operator based on a temperature term.

Example

import numpy as np
from k3_node.layers import SoftmaxAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = SoftmaxAggregation(learn=True)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

PowerMeanAggregation

`k3_node.layers.aggr.PowerMeanAggregation`

Bases: Aggregation

The powermean aggregation operator based on a power term.

Example

import numpy as np
from k3_node.layers import PowerMeanAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = PowerMeanAggregation(learn=True)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

StdAggregation

`k3_node.layers.aggr.StdAggregation`

Bases: Aggregation

An aggregation operator that takes the feature-wise standard deviation across a set of elements.

Example

import numpy as np
from k3_node.layers import StdAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = StdAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

VarAggregation

`k3_node.layers.aggr.VarAggregation`

Bases: Aggregation

An aggregation operator that takes the feature-wise variance across a set of elements.

Example

import numpy as np
from k3_node.layers import VarAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = VarAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

MedianAggregation

`k3_node.layers.aggr.MedianAggregation`

Bases: QuantileAggregation

An aggregation operator that returns the feature-wise median of a set.

Example

import numpy as np
from k3_node.layers import MedianAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MedianAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

QuantileAggregation

`k3_node.layers.aggr.QuantileAggregation`

Bases: Aggregation

An aggregation operator that returns the feature-wise :math:q-th quantile of a set :math:\mathcal{X}.

Example

import numpy as np
from k3_node.layers import QuantileAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = QuantileAggregation(q=0.75)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

DegreeScalerAggregation

`k3_node.layers.aggr.DegreeScalerAggregation`

Bases: Aggregation

Combines one or more aggregators and transforms its output with one or more scalers as introduced in the "Principal Neighbourhood Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718>_ paper.

Example

import numpy as np
from k3_node.layers import DegreeScalerAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

deg = np.array([0, 3, 5, 2])  # in-degree histogram of the training graphs
aggr = DegreeScalerAggregation(aggr=["mean", "max"], scaler=["identity", "amplification"], deg=deg)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 32): 2 aggregators x 2 scalers x 8 features

Neural & Multi-Aggregations

MultiAggregation

`k3_node.layers.aggr.MultiAggregation`

Bases: Aggregation

Performs aggregations with one or more aggregators and combines aggregated results, as described in the "Principal Neighbourhood Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718> and "Adaptive Filters and Aggregator Fusion for Efficient Graph Convolutions" <https://arxiv.org/abs/2104.01481> papers.

Example

import numpy as np
from k3_node.layers import MultiAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MultiAggregation(aggrs=["sum", "mean", "max"])
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 24)

AttentionalAggregation

`k3_node.layers.aggr.AttentionalAggregation`

Bases: Aggregation

The soft attention aggregation layer from the "Graph Matching Networks for Learning the Similarity of Graph Structured Objects" <https://arxiv.org/abs/1904.12787>_ paper.

Example

import numpy as np
import keras
from k3_node.layers import AttentionalAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = AttentionalAggregation(gate_nn=keras.layers.Dense(1), nn=keras.layers.Dense(16))
out = aggr(x, index=index, dim_size=2)  # attention-weighted sum per set
print(tuple(out.shape))  # (2, 16)

Set2Set

`k3_node.layers.aggr.Set2Set`

Bases: Aggregation

The Set2Set aggregation operator based on iterative content-based attention, as described in the "Order Matters: Sequence to sequence for Sets" <https://arxiv.org/abs/1511.06391>_ paper.

Example

import numpy as np
from k3_node.layers import Set2Set

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = Set2Set(in_channels=8, processing_steps=2)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

DeepSetsAggregation

`k3_node.layers.aggr.DeepSetsAggregation`

Bases: Aggregation

Performs Deep Sets aggregation in which the elements to aggregate are first transformed by a Multi-Layer Perceptron (MLP) :math:\phi_{\mathbf{\Theta}}, summed, and then transformed by another MLP :math:\rho_{\mathbf{\Theta}}.

Example

import numpy as np
import keras
from k3_node.layers import DeepSetsAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = DeepSetsAggregation(local_nn=keras.layers.Dense(16), global_nn=keras.layers.Dense(16))
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

MLPAggregation

`k3_node.layers.aggr.MLPAggregation`

Bases: Aggregation

Performs MLP aggregation in which the elements to aggregate are flattened into a single vectorial representation, and are then processed by a Multi-Layer Perceptron (MLP).

Example

import numpy as np
from k3_node.layers import MLPAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = MLPAggregation(in_channels=8, out_channels=16, max_num_elements=5)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

LSTMAggregation

`k3_node.layers.aggr.LSTMAggregation`

Bases: Aggregation

Performs LSTM-style aggregation in which the elements to aggregate are interpreted as a sequence, as described in the "Inductive Representation Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>_ paper.

Example

import numpy as np
from k3_node.layers import LSTMAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = LSTMAggregation(in_channels=8, out_channels=16)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

GRUAggregation

`k3_node.layers.aggr.GRUAggregation`

Bases: Aggregation

Performs GRU aggregation in which the elements to aggregate are interpreted as a sequence, as described in the "Graph Neural Networks with Adaptive Readouts" <https://arxiv.org/abs/2211.04952>_ paper.

Example

import numpy as np
from k3_node.layers import GRUAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = GRUAggregation(in_channels=8, out_channels=16)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

SetTransformerAggregation

`k3_node.layers.aggr.SetTransformerAggregation`

Bases: Aggregation

Performs "Set Transformer" aggregation in which the elements to aggregate are processed by multi-head attention blocks, as described in the "Graph Neural Networks with Adaptive Readouts" <https://arxiv.org/abs/2211.04952>_ paper.

Example

import numpy as np
from k3_node.layers import SetTransformerAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = SetTransformerAggregation(channels=8, num_seed_points=2)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

GraphMultisetTransformer

`k3_node.layers.aggr.GraphMultisetTransformer`

Bases: Aggregation

The Graph Multiset Transformer pooling operator from the "Accurate Learning of Graph Representations with Graph Multiset Pooling" <https://arxiv.org/abs/2102.11533>_ paper.

Example

import numpy as np
from k3_node.layers import GraphMultisetTransformer

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = GraphMultisetTransformer(channels=8, k=2)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

SortAggregation

`k3_node.layers.aggr.SortAggregation`

Bases: Aggregation

The pooling operator from the "An End-to-End Deep Learning Architecture for Graph Classification" <https://www.cse.wustl.edu/~muhan/papers/AAAI_2018_DGCNN.pdf>_ paper, where node features are sorted in descending order based on their last feature channel. The first :math:k nodes form the output of the layer.

Example

import numpy as np
from k3_node.layers import SortAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = SortAggregation(k=2)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

VariancePreservingAggregation

`k3_node.layers.aggr.VariancePreservingAggregation`

Bases: Aggregation

Performs the Variance Preserving Aggregation (VPA) from the "GNN-VPA: A Variance-Preserving Aggregation Strategy for Graph Neural Networks" <https://arxiv.org/abs/2403.04747>_ paper.

Example

import numpy as np
from k3_node.layers import VariancePreservingAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = VariancePreservingAggregation()
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 8)

PatchTransformerAggregation

`k3_node.layers.aggr.PatchTransformerAggregation`

Bases: Aggregation

Performs patch transformer aggregation in which the elements to aggregate are processed by multi-head attention blocks across patches.

Example

import numpy as np
from k3_node.layers import PatchTransformerAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = PatchTransformerAggregation(in_channels=8, out_channels=16, patch_size=2, hidden_channels=8)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

LCMAggregation

`k3_node.layers.aggr.LCMAggregation`

Bases: Aggregation

The Learnable Commutative Monoid aggregation from the "Learnable Commutative Monoids for Graph Neural Networks" <https://arxiv.org/abs/2212.08541>_ paper.

Example

import numpy as np
from k3_node.layers import LCMAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = LCMAggregation(in_channels=8, out_channels=16)
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

EquilibriumAggregation

`k3_node.layers.aggr.EquilibriumAggregation`

Bases: Aggregation

The equilibrium aggregation layer from the "Equilibrium Aggregation: Encoding Sets via Optimization" <https://arxiv.org/abs/2202.12795>_ paper.

Example

import numpy as np
from k3_node.layers import EquilibriumAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = EquilibriumAggregation(in_channels=8, out_channels=16, num_layers=[8])
out = aggr(x, index=index, dim_size=2)
print(tuple(out.shape))  # (2, 16)

FusedAggregation

`k3_node.layers.aggr.FusedAggregation`

Bases: Aggregation

Helper class to fuse computation of multiple aggregations together.

Example

import numpy as np
from k3_node.layers import FusedAggregation

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1

aggr = FusedAggregation(aggrs=["sum", "mean", "max"])
outs = aggr(x, index=index, dim_size=2)  # one result per aggregation, computed together
print(len(outs), tuple(outs[0].shape))  # 3 (2, 8)