GCN_reweight(in_channels, out_channels, aggr, improved=False, cached=False, add_self_loops=False, normalize=True, bias=True, **kwargs)

Bases: MessagePassing

Graph Convolutional Network with edge reweighting mechanism.

Parameters:
  • in_channels (int) –

    Input feature dimensionality

  • out_channels (int) –

    Output feature dimensionality

  • aggr (str) –

    Aggregation method ('add', 'mean', etc.)

  • improved (bool, default: False ) –

    If True, use improved GCN normalization. Default: False

  • cached (bool, default: False ) –

    Whether to cache normalized adjacency matrix. Default: False

  • add_self_loops (bool, default: False ) –

    Whether to add self-loops. Default: False

  • normalize (bool, default: True ) –

    Whether to apply normalization. Default: True for non-add aggregation

  • bias (bool, default: True ) –

    Whether to include bias. Default: True
forward(x, edge_index, edge_weight, lmda)

Forward pass of the reweighted GCN layer.

Parameters:
  • x (Tensor) –

    Node feature matrix [num_nodes, in_channels]

  • edge_index (Tensor or SparseTensor) –

    Graph connectivity in COO format [2, num_edges] or sparse format

  • edge_weight (Tensor) –

    Edge weights for reweighting [num_edges]

  • lmda (float) –

    Interpolation factor between normalized and reweighted adjacency
Returns:
  • Tensor –

    Updated node features [num_nodes, out_channels]
Notes

  • Edge Weight Processing:

    • Store original weights as reweighting values
    • Initialize base weights as ones

  • Normalization (if enabled):

    • Compute symmetric normalization if not cached
    • Use cached values if available
    • Handle both dense and sparse formats

  • Feature Transformation:

    • Linear projection of input features
    • Message passing with interpolated weights
    • Optional bias addition
message(x_j, edge_index, edge_weight, edge_rw, lmda)

Compute messages from source nodes with edge reweighting.

Parameters:
  • x_j (Tensor) –

    Source node features [num_edges, out_channels]

  • edge_index (Tensor) –

    Edge indices [2, num_edges]

  • edge_weight (Tensor) –

    Normalized edge weights [num_edges]

  • edge_rw (Tensor) –

    Original edge weights for reweighting [num_edges]

  • lmda (float) –

    Interpolation factor between normalized and reweighted messages
Returns:
  • Tensor –

    Computed messages with interpolated weights [num_edges, out_channels]
Notes
  • Computes messages with interpolated weights
  • Handles both normalized and original edge weights
message_and_aggregate(adj_t, x)

Fused message computation and aggregation for sparse tensors.

Parameters:
  • adj_t (SparseTensor) –

    Transposed adjacency matrix in sparse format

  • x (Tensor) –

    Node feature matrix [num_nodes, out_channels]
Returns:
  • Tensor –

    Aggregated messages [num_nodes, out_channels]
Notes
  • Optimized sparse matrix multiplication
  • Uses specified aggregation method (sum, mean, etc.)
  • More efficient than separate message and aggregation
GS_reweight(in_channels, out_channels, reducer, normalize_embedding=False)

Bases: MessagePassing

GraphSAGE with edge reweighting mechanism.

Parameters:
  • in_channels (int) –

    Input feature dimensionality

  • out_channels (int) –

    Output feature dimensionality

  • reducer (str) –

    Aggregation method

  • normalize_embedding (bool, default: False ) –

    Whether to normalize output embeddings. Default: False
Notes
  • Two-layer transformation (neighbor + self)
  • Edge weight interpolation
  • Optional embedding normalization
forward(x, edge_index, edge_weight, lmda)

Forward pass of the reweighted GraphSAGE layer.

Parameters:
  • x (Tensor) –

    Node feature matrix [num_nodes, in_channels]

  • edge_index (Tensor) –

    Graph connectivity in COO format [2, num_edges]

  • edge_weight (Tensor) –

    Edge weights for reweighting [num_edges]

  • lmda (float) –

    Interpolation factor for edge weight importance
Returns:
  • Tensor –

    Updated node features [num_nodes, out_channels]
Notes

Initiates message passing with proper graph size information

message(x_j, edge_index, edge_weight, lmda)

Compute messages from source nodes with edge reweighting.

Parameters:
  • x_j (Tensor) –

    Source node features [num_edges, in_channels]

  • edge_index (Tensor) –

    Edge indices [2, num_edges]

  • edge_weight (Tensor) –

    Edge weights [num_edges]

  • lmda (float) –

    Interpolation factor between base and weighted features
Returns:
  • Tensor –

    Transformed and reweighted messages [num_edges, out_channels]
Notes

  • Linear transformation of source features

  • Interpolation between:

    • Transformed features (weight: 1-λ)
    • Edge-weighted features (weight: λ)
update(aggr_out, x)

Update node embeddings using aggregated messages and self-features.

Parameters:
  • aggr_out (Tensor) –

    Aggregated messages [num_nodes, out_channels]

  • x (Tensor) –

    Original node features [num_nodes, in_channels]
Returns:
  • Tensor –

    Updated node embeddings [num_nodes, out_channels]
Notes
  • Concatenate aggregated messages with self-features
  • Apply learnable transformation
  • Non-linear activation (ReLU)
  • Optional L2 normalization
ReweightGNN(input_dim, gnn_dim, output_dim, cls_dim, gnn_layers=3, cls_layers=2, backbone='GS', pooling='mean', dropout=0.5, bn=False, rw_lmda=1.0, **kwargs)

Bases: Module

Multi-layer GNN with edge reweighting mechanism.

Parameters:
  • input_dim (int) –

    Input feature dimensionality

  • gnn_dim (int) –

    Hidden dimension of GNN layers

  • output_dim (int) –

    Output dimension

  • cls_dim (int) –

    Hidden dimension of classification layers

  • gnn_layers (int, default: 3 ) –

    Number of GNN layers. Default: 3

  • cls_layers (int, default: 2 ) –

    Number of classification layers. Default: 2

  • backbone (str, default: 'GS' ) –

    GNN architecture ('GCN' or 'GS'). Default: 'GS'

  • pooling (str, default: 'mean' ) –

    Pooling method. Default: 'mean'

  • dropout (float, default: 0.5 ) –

    Dropout rate. Default: 0.5

  • bn (bool, default: False ) –

    Whether to use batch normalization. Default: False

  • rw_lmda (float, default: 1.0 ) –

    Edge reweighting interpolation factor. Default: 1.0
Notes
  • Multiple GNN layers with reweighting
  • Optional batch normalization
  • Multi-layer classification head
  • Dropout regularization
forward(data, h)

Forward pass of the network.

Parameters:
  • data (Data) –

    Graph data object containing edge indices and weights

  • h (Tensor) –

    Input node features
Returns:
  • tuple[Tensor, Tensor] –
    • Node embeddings after GNN layers
    • Classification logits
Notes
  • GNN propagation with reweighting
  • Activation and dropout
  • Classification MLP
  • Optional batch normalization
spmm(src, other, reduce='sum')

Sparse matrix multiplication with dense matrix.

Parameters:
  • src (Tensor or SparseTensor) –

    Input sparse matrix

  • other (Tensor) –

    Input dense matrix

  • reduce (str, default: 'sum' ) –

    Reduction operation ('sum', 'mean', 'min', 'max'). Default: 'sum'
Returns:
  • Tensor –

    Result of sparse-dense matrix multiplication
Notes

Supports different sparse formats and reduction operations