Skip to content

A differentiable JIT compiler for graph message passing.

Define a graph, write what each edge sends, and choose how messages combine. Call the program to obtain node results and derive gradients without a handwritten backward function. Ordinary use does not require compiler IR.

Directly supports PyTorch torch.Tensor inputs and outputs without additional wrappers. Compute gradients with torch.autograd.grad or .backward() and integrate Tiga into existing PyTorch models and training workflows.

Start here

Develop the compiler

  1. Understand the internal layers.
  2. Follow one real program through its IR.
  3. Build, test and locate the implementation.

Start after the Python programming model. The walkthrough explains notation before showing passes and generated output.

What a program looks like

This example aggregates weighted neighbor values on an NVIDIA GPU and computes gradients with PyTorch:

import torch
import tiga as tg

device = torch.device("cuda")

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

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

graph = tg.Graph.from_csr(
    torch.tensor([0, 2, 3, 5], dtype=torch.int64, device=device),
    torch.tensor([0, 2, 1, 0, 1], dtype=torch.int64, device=device),
    num_src=3,
)
x = torch.tensor([1.0, 2.0, 3.0], device=device, requires_grad=True)
weight = torch.tensor([2.0, 3.0, 4.0, 5.0, 6.0], device=device, requires_grad=True)
program = WeightedAggregation()
out = program(
    graph=graph,
    src={"x": x},
    dst={},
    edge={"weight": weight},
)
print(out.tolist())  # [11.0, 8.0, 17.0]
dx, dw = torch.autograd.grad(out.sum(), (x, weight))
print(dx.tolist())   # [7.0, 10.0, 3.0]
print(dw.tolist())   # [1.0, 3.0, 2.0, 1.0, 2.0]

The graph supplies the edges; edge() supplies each message; the reducer tg.sum() combines messages per destination node. tg is only an import alias for tiga, not another package.

Graph indices, inputs, outputs and gradients reside on the GPU, with no tensor wrapping or change of autograd interface. See the installation guide for environment setup. Execution diagnostics explains how to inspect the selected path.

Current scope

Supports CPU and NVIDIA GPU execution, with hierarchical storage and distributed execution for large graphs. Standard computations accept torch.Tensor directly; advanced runtime features such as hierarchical storage and distributed execution use native tg.Tensor, as described in memory and distributed execution. See the support matrix for operation, dtype and execution-path coverage.

Sparse graphs and neural network operators

Weighted neighbor aggregation is a sparse matrix operation; custom edge functions extend it to more general graph programs. The PyTorch sparse and message passing guide connects torch.sparse.mm, CSR graphs and graph neural network (GNN) aggregation with one checked example. Tiga compiles these operators; it is not a replacement for a complete model and training framework.

Understand the implementation

The compiler turns a graph computation into a traversal and an execution plan. These are internal decisions, not extra APIs to write. The compiler introduction explains their purpose; the IR walkthrough shows actual gf.apply, gf_iter.traverse, gf_kernel.launch and task operations. The technical report develops the programming model and compiler design, with evaluation and reproduction data.

Evaluate performance

The performance and scalability report compares execution time, memory use, large-graph capacity and distributed execution for specified workloads. The measurement protocol defines timing and validation; the benchmark archive contains additional measured cases.

For signatures, use the Python API reference. For project positioning, see comparisons.