PyTorch sparse matrix multiplication and graph message passing

PyTorch exposes sparse matrix operations through torch.sparse. Tiga expresses computation as message passing along graph edges. For weighted neighbor summation, these are two ways to express the same operation. General edge functions need not be matrix multiplication.

How a sparse matrix represents a graph

In compressed sparse row (CSR) format, a row represents a destination node. Each stored column index identifies a source node, and the corresponding value is an edge weight. For an edge from source j to destination i:

\[ y_i = \sum_{e=(j\to i)} w_e x_j. \]

With feature vectors, the same expression is sparse matrix–dense matrix multiplication (SpMM). torch.sparse.mm(adjacency, features) and a Tiga edge.weight * src.x message with a sum reducer describe this weighted sum. The graph programming model illustrates how graph connections map to data; the example below fixes the row/column convention explicitly.

A GPU example with forward and gradient checks

Complete installation first, including a CUDA-compatible Torch installation. This example compares outputs and gradients with independent Torch inputs; it is a correctness check, not a benchmark.

Run python examples/pytorch_sparse_message_passing.py --device cuda. The same source accepts --device cpu for a CPU check.

"""Compare weighted graph aggregation with torch.sparse.mm, including gradients."""
import argparse

import torch
import tiga as tg


class WeightedAggregation(tg.MessagePassing):
    reducer = tg.sum()

    def edge(self, src, dst, edge):
        return edge.weight * src.x


def run(device="cuda"):
    # CSR rows are destinations; column indices identify source nodes.
    row = torch.tensor([0, 2, 3, 5], dtype=torch.int64, device=device)
    col = torch.tensor([0, 2, 1, 0, 1], dtype=torch.int64, device=device)
    x = torch.tensor([1., 2., 3.], device=device, requires_grad=True)
    weight = torch.tensor([2., 3., 4., 5., 6.], device=device,
                          requires_grad=True)

    graph = tg.Graph.from_csr(row, col, num_src=3)
    result = WeightedAggregation()(
        graph=graph, src={"x": x}, dst={}, edge={"weight": weight})

    # Independent leaves make the two gradient computations independent.
    x_ref = x.detach().clone().requires_grad_()
    weight_ref = weight.detach().clone().requires_grad_()
    adjacency = torch.sparse_csr_tensor(
        row, col, weight_ref, size=(3, 3), device=device, check_invariants=True)
    reference = torch.sparse.mm(adjacency, x_ref[:, None]).squeeze(1)

    torch.testing.assert_close(result, reference)
    gradients = torch.autograd.grad(result.sum(), (x, weight))
    reference_gradients = torch.autograd.grad(
        reference.sum(), (x_ref, weight_ref))
    for actual, expected in zip(gradients, reference_gradients):
        torch.testing.assert_close(actual, expected)
    print("Output:", result.detach().tolist())  # [11.0, 8.0, 17.0]
    print("Node gradients:", gradients[0].tolist())  # [7.0, 10.0, 3.0]
    print("Edge gradients:", gradients[1].tolist())  # [1., 3., 2., 1., 2.]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
    run(parser.parse_args().device)

Both implementations return ordinary torch.Tensor values. The five stored edges remain separate edge weights. Rows without incoming edges sum to zero; repeated source–destination pairs contribute once per stored edge in this sum. Other reducers have their own empty-neighborhood rules. Torch sparse layout, device and gradient coverage are defined by torch.sparse.mm; Tiga coverage is listed in the support matrix.

Where graph neural networks fit

A graph neural network (GNN) combines graph aggregation with feature transformations and nonlinearities. The sum above is an aggregation building block, not a complete normalized graph convolutional network (GCN) layer. Self-loops, degree normalization, learned transformations and training objectives must be specified by the model. See the multi-feature GCN-shaped aggregation example and Torch model integration.

PyTorch Geometric (PyG) provides graph learning layers and training utilities. Tiga instead focuses on compiling graph programs, including custom edge messages and their supported gradients. It is not a drop-in replacement for every PyG layer. The comparison guide explains these abstraction boundaries.

Beyond a fixed sparse adjacency matrix

Radius and k-nearest-neighbor graphs derive neighbors from coordinates or features; stencils derive neighbors from grid rules. The dynamic relation examples explain these constructors, distance metrics and boundaries. Whether a path materializes edges or generates candidates during execution depends on the relation and selected backend.

For performance, compare the same graph, features, dtype, device and gradients. Report graph construction, compilation, warmed execution and memory separately. The performance report provides measured sparse aggregation, radius/kNN, large-graph and distributed cases; it does not establish that Tiga is faster for every graph neural network or torch.sparse operation.