Message passing and graph algorithms

Inputs and outputs use Torch tensors with torch.autograd.grad by default. The custom reducer and PageRank control-flow probe still require native tensors because their Torch execution paths are not implemented. The measured compilation artifacts below are historical native-path captures, not artifacts guaranteed by the current Torch examples; inspect kernel.explain() for the actual plan.

Static-topology MessagePassing programs and a fixed-iteration graph algorithm probe. The compiler generates the relation traversal and its VJP; the user writes only the UDFs. The message passing guide explains the interface. Each section shows the minimal core snippet; the complete runnable program is one click away in the collapsed source block.

• python examples/message_passing_autograd.py
• python examples/gcn.py
• python examples/diffusion.py
• python examples/custom_reducer.py
• python examples/compiler_probes/pagerank.py

GCN aggregation

What it is. A graph convolutional network (GCN) layer: every node of a graph carries a feature vector (here F = 2 numbers), and the layer replaces each node's features with the sum of the feature vectors of the nodes that link to it, each scaled by one scalar weight per edge. In linear-algebra terms this is a weighted sparse matrix–matrix multiplication (SpMM): the weighted adjacency matrix times the feature matrix. The example also differentiates the sum of all outputs back onto the features and the weights.

\[ \text{out}_{i,f} \;=\; \sum_{e\,=\,(j \to i)} w_e \cdot x_{j,f} \]

Torch multi-feature aggregation with edge broadcasting and automatic gradients, matching the GCN tensor shape. The GCN layer is user code, not a library operator.

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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): weight[e] * x[j, :] — (F,) per edge
        return edge.weight * src.x


output = WeightedFeatureAggregation()(
    graph=graph, src={"x": x}, dst={}, edge={"weight": weight})  # (N, F)

Full source: examples/gcn.py (runs as-is)

"""feature MessagePassing and automatic gradients (SpMM shape)."""

import torch
import tiga as tg


# 3 nodes, 5 edges: 0→0, 2→0, 1→1, 0→2, 1→2
row_ptr = torch.tensor([0, 2, 3, 5], dtype=torch.int64)        # (N+1,)
col_idx = torch.tensor([0, 2, 1, 0, 1], dtype=torch.int64)     # (E,)
graph = tg.Graph.from_csr(row_ptr, col_idx, num_src=3, validate="full")

x = torch.tensor(
    [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], requires_grad=True)        # (N, F)
# [E, 1] broadcasts over the feature axis without changing graph semantics.
weight = torch.tensor(
    [[0.5], [1.0], [2.0], [1.5], [0.25]], requires_grad=True)        # (E, 1)


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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): weight[e] * x[j, :] — (F,) per edge
        return edge.weight * src.x


output = WeightedFeatureAggregation()(
    graph=graph, src={"x": x}, dst={}, edge={"weight": weight})  # (N, F)

dx, dweight = torch.autograd.grad(output.sum(), (x, weight))
# d x[j, :] = Σ weight[e] over edges e out of j (broadcast over F)
# d weight[e] = Σ_f x[src(e), f]

Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)

CPUCUDA

### WeightedFeatureAggregation
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

### WeightedFeatureAggregation
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

Directional diffusion

What it is. Diffusion: a quantity stored on the nodes of a graph (heat, concentration) flows along the edges from larger values toward smaller ones until it evens out. The flux into a node is the sum over its incoming edges of the edge's conductivity times the difference between the neighbor's value and the node's own. The example advances one explicit Euler step — each node moves by a small time step Δt times its net flux — on a 4-node ring where every node exchanges with two neighbors.

\[ \text{flux}_i \;=\; \sum_{e\,=\,(j \to i)} c_e \, (u_j - u_i), \qquad u^{+}_i \;=\; u_i + \Delta t \cdot \text{flux}_i \]

Both the field u and the edge conductivities receive compiler-generated gradients.

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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): conductivity[e] * (u[j] − u[i]) — scalar
        return edge.conductivity * (src.u + (-1.0) * dst.u)

    def node(self, dst, flux, dt):
        # flux: (N,) reduced messages; dt arrives from the call site
        return dst.u + dt * flux


next_u = Diffusion()(
    graph=graph,
    ndata={"u": u},  # homogeneous graph: one declaration, both roles
    edge={"conductivity": conductivity},
    dt=0.1,
)  # (N,)

Full source: examples/diffusion.py (runs as-is)

"""differentiable diffusion UDF on an explicit CSR stencil."""

import torch
import tiga as tg


# 4 nodes, 8 edges: dst i gathers from two sources each
nodes = 4
row_ptr = torch.tensor([0, 2, 4, 6, 8], dtype=torch.int64)    # (N+1,)
col_idx = torch.tensor([3, 1, 0, 2, 1, 3, 2, 0], dtype=torch.int64)  # (E,)
graph = tg.Graph.from_csr(row_ptr, col_idx, num_src=nodes, validate="full")

u = torch.tensor([0.0, 1.0, 0.5, -0.5], requires_grad=True)         # (N,)
conductivity = torch.tensor([1.0] * 8, requires_grad=True)          # (E,)


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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): conductivity[e] * (u[j] − u[i]) — scalar
        return edge.conductivity * (src.u + (-1.0) * dst.u)

    def node(self, dst, flux, dt):
        # flux: (N,) reduced messages; dt arrives from the call site
        return dst.u + dt * flux


next_u = Diffusion()(
    graph=graph,
    ndata={"u": u},  # homogeneous graph: one declaration, both roles
    edge={"conductivity": conductivity},
    dt=0.1,
)  # (N,)
# flux = [0.5, -1.5, -0.5, 1.5]

du, d_conductivity = torch.autograd.grad(next_u.sum(), (u, conductivity))
# every node has in-degree = out-degree = 2, so the flux terms cancel in du
# d conductivity[e] = dt · (u[src(e)] − u[dst(e)])

Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)

CPUCUDA

### Diffusion
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

### Diffusion
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

Differentiable MessagePassing UDF

What it is. This example treats nodes as points carrying a temperature and a bias, and edges as conductors: each edge forwards its conductivity times the temperature of its source node to its destination, which adds its own bias. The example then asks for the derivative of the total output with respect to every input, delivered as a vector–Jacobian product (VJP) — the adjoint of the forward computation — without any hand-written backward code.

The example computes a conductive aggregation with a destination bias,

\[ \text{out}_i \;=\; b_i \;+\; \sum_{e\,=\,(j \to i)} c_e \cdot T_j \]

and differentiates sum(out) with respect to all three Torch inputs using torch.autograd.grad. Native-only checkpoint policy inspection is separate: python examples/native_checkpoint.py.

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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): conductivity[e] * temperature[j] — scalar
        return edge.conductivity * src.temperature

    def node(self, dst, incoming):
        # incoming: (N,) reduced messages; dst.bias: (N,)
        return incoming + dst.bias


kernel = ConductiveAggregation()
output = kernel(  # (N,)
    graph=graph,
    src={"temperature": temperature},
    dst={"bias": bias},
    edge={"conductivity": conductivity},
)

Full source: examples/message_passing_autograd.py (runs as-is)

"""Differentiable MessagePassing UDF with no user-written backward."""

import torch
import tiga as tg


# 3 nodes, 5 edges: 0→0, 2→0, 1→1, 0→2, 1→2
row_ptr = torch.tensor([0, 2, 3, 5], dtype=torch.int64)        # (N+1,)
col_idx = torch.tensor([0, 2, 1, 0, 1], dtype=torch.int64)     # (E,)
graph = tg.Graph.from_csr(row_ptr, col_idx, num_src=3, validate="full")

temperature = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)             # (N,)
conductivity = torch.tensor([2.0, 3.0, 4.0, 5.0, 6.0], requires_grad=True)  # (E,)
bias = torch.tensor([0.1, 0.2, 0.3], requires_grad=True)                    # (N,)


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

    def edge(self, src, dst, edge):
        # message on edge e=(j→i): conductivity[e] * temperature[j] — scalar
        return edge.conductivity * src.temperature

    def node(self, dst, incoming):
        # incoming: (N,) reduced messages; dst.bias: (N,)
        return incoming + dst.bias


kernel = ConductiveAggregation()
output = kernel(  # (N,)
    graph=graph,
    src={"temperature": temperature},
    dst={"bias": bias},
    edge={"conductivity": conductivity},
)
# out = [2·1 + 3·3 + 0.1,  4·2 + 0.2,  5·1 + 6·2 + 0.3]

loss = output.sum()                                                      # scalar
d_temperature, d_conductivity, d_bias = torch.autograd.grad(
    loss, (temperature, conductivity, bias))
# d temperature[j] = Σ conductivity[e] over edges e out of j
# d conductivity[e] = temperature[src(e)]

# Native-only compiler checkpoint inspection is a separate advanced example:
# examples/native_checkpoint.py

if __name__ == "__main__":
    print(output.tolist())
    print(d_temperature.tolist())

Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)

CPUCUDA

### ConductiveAggregation
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

### ConductiveAggregation
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

User-defined reducer

What it is. How to compute a mean when the only thing the machine can do in parallel is combine values pairwise. The trick is to carry a pair — running sum and item count — instead of a single number: each value v becomes (v, 1), pairs merge by adding both components, and a final division turns the accumulated pair into the mean. Because merging is associative and commutative, the combine order — and hence the parallel schedule — is free.

\[ \text{lift}(v) = (v,\, 1) \qquad (s_l, n_l) \oplus (s_r, n_r) = (s_l + s_r,\; n_l + n_r) \qquad \text{mean}_i = \frac{s_i}{n_i} \]

Where each method runs — the same computation as an ordinary Python loop:

for i in range(num_dst):                 # per destination node
    state = identity()                   # (0.0, 0.0) — before its edges
    for e in edges_into(i):
        message = edge(src, dst, edge)   # the edge() UDF, here src.value
        state = combine(state, lift(message))
    out[i] = finalize(state)             # s / n — after its edges

# node 0 traces:  (0.0, 0.0) → combine((0.0,0.0), lift(1.0)) = (1.0, 1)
#                          → combine((1.0,1),   lift(4.0))   = (5.0, 2)
#                          → finalize((5.0, 2)) = 2.5

The loop above shows the logical order, a left fold. Because the algebra is declared associative and commutative, the compiler may combine in any order or as a tree — that is what makes the parallel schedule free. It proves the tuple state componentwise additive, lowers the four regions, and generates the backward through the mean as well.

class Mean(tg.Reducer):
    associative = True
    commutative = True

    def identity(self):
        return 0.0, 0.0  # (sum, count)

    def lift(self, value):
        return value, 1.0

    def combine(self, left, right):
        return left[0] + right[0], left[1] + right[1]

    def finalize(self, state):
        return state[0] / state[1]


class NeighborMean(tg.MessagePassing):
    reducer = Mean()

    def edge(self, src, dst, edge):
        # staged reducer item: (src.value[j], 1) per edge
        return self.reducer(src.value)


kernel = NeighborMean()
output = kernel(graph=graph, src={"value": values}, dst={})  # (N,)

Full source: examples/custom_reducer.py (runs as-is)

"""Advanced custom reducer: this algebra currently needs the native runtime."""

import tiga as tg


# 3 nodes, 5 edges: 0→0, 2→0, 1→1, 0→2, 1→2
row_ptr = tg.tensor([0, 2, 3, 5], dtype=tg.int64)    # (N+1,)
col_idx = tg.tensor([0, 2, 1, 0, 1], dtype=tg.int64)  # (E,)
graph = tg.Graph.from_csr(row_ptr, col_idx, num_src=3, validate="full")
values = tg.tensor([1.0, 2.0, 4.0], requires_grad=True)          # (N,)


class Mean(tg.Reducer):
    associative = True
    commutative = True

    def identity(self):
        return 0.0, 0.0  # (sum, count)

    def lift(self, value):
        return value, 1.0

    def combine(self, left, right):
        return left[0] + right[0], left[1] + right[1]

    def finalize(self, state):
        return state[0] / state[1]


class NeighborMean(tg.MessagePassing):
    reducer = Mean()

    def edge(self, src, dst, edge):
        # staged reducer item: (src.value[j], 1) per edge
        return self.reducer(src.value)


kernel = NeighborMean()
output = kernel(graph=graph, src={"value": values}, dst={})  # (N,)

gradient = tg.autograd.grad(output.sum(), values)
# d values[j] = Σ 1/degree(dst(e)) over edges e out of j

# Inspect the selected plan: kernel.explain()


def main() -> None:
    return output, gradient


if __name__ == "__main__":
    main()

Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)

CPUCUDA

### NeighborMean
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: proved-componentwise-additive-udf
variant cache: hits=0, misses=1

### NeighborMean
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: proved-componentwise-additive-udf
variant cache: hits=0, misses=1

Fixed-iteration PageRank probe

What it is. PageRank: the algorithm that ranks web pages by importance. Every page starts with equal rank; in each round a page divides its rank equally among its outgoing links and collects the rank arriving through its in-links, blended with a small constant (1 − damping)/N that models jumping to a random page. A page with no outgoing links (a dangling node) would leak rank out of the system, so its rank is redistributed uniformly to all pages. Here: 4 pages, 3 links, damping 0.85, 20 rounds.

\[ \text{base} = \frac{1-d}{N} + \frac{d}{N} \sum_{i\ \mathrm{dangling}} \text{rank}_i, \qquad \text{rank}'_i = \text{base} + d \sum_{e\,=\,(j \to i)} \frac{\text{rank}_j}{\mathrm{deg}^{\text{out}}_j} \]

The iteration is a plain Python for loop under @tg.jit; 20 iterations are captured once as a single gf_control.repeat op — the loop stays rolled, no unrolling (the tests verify this IR shape and the rank values against a framework-independent reference). It is a compiler probe, not a workload-specific core operator or a performance claim.

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

        def edge(self, src, dst, edge):
            return src.rank * edge.inverse_out_degree

        def node(self, dst, incoming, damping, base):
            return base + damping * incoming

    step = PageRankStep()

    @tg.jit
    def iterate(rank, inverse_out_degree, dangling, damping, iterations):
        for _ in range(iterations):
            dangling_mass = (rank * dangling).sum()
            base = (1.0 - damping) / 4.0 + damping * dangling_mass / 4.0
            rank = step(
                graph=graph,
                src={"rank": rank},
                dst={},
                edge={"inverse_out_degree": inverse_out_degree},
                damping=damping,
                base=base,
            )
        return rank

    rank = iterate(rank, inverse_out_degree, dangling, damping, iterations)

Full source: examples/compiler_probes/pagerank.py (runs as-is)

"""PageRank as a compiler probe, not a core Tiga algorithm."""

from __future__ import annotations

import tiga as tg


def run(iterations: int = 20, damping: float = 0.85):
    # Incoming CSR for 0 <- 2, 1 <- 0, 2 <- 1 and dangling vertex 3.
    row_ptr = tg.tensor([0, 1, 2, 3, 3], dtype=tg.int64)      # (N+1,)
    col_idx = tg.tensor([2, 0, 1], dtype=tg.int64)            # (E,)
    graph = tg.Graph.from_csr(
        row_ptr, col_idx, num_src=4, validate="full")

    rank = tg.tensor([0.25] * 4, dtype=tg.float32)            # (N,)
    inverse_out_degree = tg.tensor([1.0, 1.0, 1.0], dtype=tg.float32)  # (E,)
    dangling = tg.tensor([0.0, 0.0, 0.0, 1.0], dtype=tg.float32)       # (N,)

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

        def edge(self, src, dst, edge):
            return src.rank * edge.inverse_out_degree

        def node(self, dst, incoming, damping, base):
            return base + damping * incoming

    step = PageRankStep()

    @tg.jit
    def iterate(rank, inverse_out_degree, dangling, damping, iterations):
        for _ in range(iterations):
            dangling_mass = (rank * dangling).sum()
            base = (1.0 - damping) / 4.0 + damping * dangling_mass / 4.0
            rank = step(
                graph=graph,
                src={"rank": rank},
                dst={},
                edge={"inverse_out_degree": inverse_out_degree},
                damping=damping,
                base=base,
            )
        return rank

    rank = iterate(rank, inverse_out_degree, dangling, damping, iterations)

    canonical_ir = rank.mlir()
    actual = rank.tolist()

    expression_nodes = sum(
        line.startswith("%") for line in rank.expression().splitlines())
    return actual, expression_nodes, canonical_ir


if __name__ == "__main__":
    run()

# Inspect: rank.expression(), rank.mlir()  (via run()'s return values)

Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)

CPUCUDA

### PageRankStep
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1

### PageRankStep
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1