Dynamic and generated relations

Relations can be generated from coordinates or from fixed rules. “Generated” does not necessarily mean the topology changes on every call. These examples accept torch.Tensor directly, with no additional wrappers.

Relation	What determines neighbors	Typical use
Graph.stencil	Grid dimensions and fixed offsets	Image filters and local grid updates
Graph.radius	Distances and a cutoff	Particles and spatial neighborhoods
Graph.knn	Distance ranking and k	Fixed-degree point-cloud computation

Five-point stencil on a regular grid

A stencil specifies fixed neighbor offsets. This example smooths a 2-D grid by averaging the center, up, down, left and right values. No positions array or nearest-neighbor search is needed.

dims=(5, 5) specifies five rows and five columns; offsets are (row, column). Nodes use row-major ordering: coordinate (2, 2) is node 12, with neighbors 7, 17, 11, 13. The tg.stencil.von_neumann() macro includes the center by default.

Neighborhood macros

Von Neumann uses Manhattan distance; Moore uses Chebyshev distance. Both expand integer offsets with distance at most radius, infer the dimension from dims, and include the center by default. For integer Manhattan distance, <= 1 is equivalent to < 2.

Macro	2-D offsets	Number including center
tg.stencil.von_neumann()	center and four axis neighbors	5
tg.stencil.von_neumann(2)	diamond of radius 2	13
tg.stencil.moore()	full 3×3 neighborhood	9

graph = tg.Graph.stencil((64, 64), device="cuda")  # Default: von_neumann().
graph = tg.Graph.stencil((64, 64), tg.stencil.moore(), device="cuda")
four_neighbors = tg.stencil.von_neumann(include_center=False)
offsets = four_neighbors.offsets(2)  # ((-1, 0), (0, -1), (0, 1), (1, 0))

Explicit offset tuples remain supported. Macros expand in lexicographic order; edge fields must follow that order after boundary filtering. On a small periodic grid, different offsets can wrap to the same source: those edges remain distinct and contribute separately. See parameter reference.

Five-point averaging

\[ y_{r,c}=\frac{x_{r,c}+x_{r-1,c}+x_{r+1,c}+x_{r,c-1}+x_{r,c+1}}{5} \]

import torch
import tiga as tg


class GridAverage(tg.MessagePassing):
    reducer = tg.mean()

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


def run(device="cuda"):
    graph = tg.Graph.stencil(
        (5, 5),
        tg.stencil.von_neumann(),  # Radius 1, including the center: five points.
        periodic=True, device=device,
    )
    x = torch.arange(25, dtype=torch.float32, device=device).requires_grad_()
    y = GridAverage()(graph=graph, src={"x": x}, dst={}, edge={})
    dx, = torch.autograd.grad(y.sum(), (x,))
    return y.reshape(5, 5), dx.reshape(5, 5)

run("cuda") returns the output and gradient, both shaped (5, 5). The center output is 12.0 and the top-left output is 6.0; the gradient of the total output sum is 1.0 for every input.

• periodic=True wraps across boundaries, so the top-left node still has five inputs.
• periodic=False omits out-of-grid neighbors; it does not zero-pad. With tg.mean(), a corner averages only its three valid inputs.
• Zero/constant padding is not implemented by this constructor. Omitting an edge is not equivalent to sending a zero-valued message, especially for mean.
• The current Graph.stencil constructor generates CSR. A rule-based definition does not imply edge-free execution. Reuse the graph while dimensions and offsets remain unchanged.

Run python examples/stencil_message_passing.py. CUDA is the default; add --device cpu for a CPU-only environment.

Full source: examples/stencil_message_passing.py

"""Five-point periodic grid averaging with Torch tensors and autograd."""
import argparse

import torch
import tiga as tg


class GridAverage(tg.MessagePassing):
    reducer = tg.mean()

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


def run(device="cuda"):
    graph = tg.Graph.stencil(
        (5, 5),
        tg.stencil.von_neumann(),  # Radius 1, including the center: five points.
        periodic=True, device=device,
    )
    x = torch.arange(25, dtype=torch.float32, device=device).requires_grad_()
    y = GridAverage()(graph=graph, src={"x": x}, dst={}, edge={})
    dx, = torch.autograd.grad(y.sum(), (x,))
    return y.reshape(5, 5), dx.reshape(5, 5)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
    args = parser.parse_args()
    result, gradient = run(args.device)
    print(result.tolist())
    print(gradient.tolist())  # Every input contributes a total weight of 1.

The radius and kNN examples below instead depend on positions. Collapsed compilation records describe particular captured runs; kernel.explain() reports the actual execution plan.

Euclidean and custom distances

Euclidean distance is the default for both constructors. Cosine distance is one minus cosine similarity; smaller distances mean closer neighbors.

Constructor	Current distance support	Selection
Graph.radius(x, cutoff=r)	Euclidean by default; custom metric callable	distance ≤ cutoff
Graph.knn(x, k=k)	Euclidean only; no metric argument yet	k nearest candidates

import torch
import torch.nn.functional as F
import tiga as tg


def cosine_distance(src, dst, edge):
    # One scalar per candidate pair. Inputs below have nonzero norms.
    return (1 - F.cosine_similarity(src.position, dst.position, dim=-1)).clamp_min(0)


def relations(device="cuda"):
    x = torch.tensor([[1., 0.], [.8, .6], [0., 1.], [-1., 0.]], device=device)
    euclidean_radius = tg.Graph.radius(x, cutoff=0.7)
    euclidean_knn = tg.Graph.knn(x, k=1)
    cosine_radius = tg.Graph.radius(x, cutoff=0.25, metric=cosine_distance)

    # On unit vectors, Euclidean and cosine distances have the same ranking.
    unit_x = F.normalize(x, dim=-1)
    cosine_knn = tg.Graph.knn(unit_x, k=1)
    return euclidean_radius, euclidean_knn, cosine_radius, cosine_knn

metric(src, dst, edge) receives batches of candidate pairs. src.position and dst.position have shape (P, D); extra fields appear on the endpoint namespaces. Return a floating tensor of shape (P,) on the same device. The metric callback receives edge.displacement (minimum-image displacement when periodic); edge.distance is its result, not an input to the callback. An optional select(src, dst, edge) runs afterward and can read edge.distance and edge.cutoff; it returns a boolean (P,) tensor and only removes edges. The example uses Torch callbacks with Torch coordinates; Torch-free native callbacks must use native Tensor operations and currently require CPU.

The cosine radius example selects similarity ≥ 0.75, excludes self edges, and connects only nodes 0 and 1 in both directions. Arbitrary custom metrics use an all-pairs correctness path: the Euclidean cell-list bound no longer applies. Even a callable that recomputes Euclidean distance takes this path; omit metric to retain the built-in optimizations. No arbitrary-metric fused performance claim is implied.

For cosine kNN, normalize nonzero vectors first:

\[ \|\hat{x}-\hat{y}\|_2^2 = 2\bigl(1-\cos(x,y)\bigr) \]

The rankings agree, apart from numerical ties. Normalize both queries and candidates for bipartite kNN. This changes the coordinates used by the graph; it does not redefine any reported Euclidean distance as cosine distance. Reject zero/near-zero vectors or define their behavior explicitly before using this equivalence. Discrete neighbor selection is not differentiable; gradients through supported selected-edge computations hold the neighbor set fixed.

Run python examples/distance_metrics.py on CUDA, or append --device cpu. The tiny example explicitly materializes CSR only to print and inspect the chosen neighbors.

Differentiable radius relation

What it is. A radius graph is built from point positions: two points are connected whenever their Euclidean distance is at most a cutoff r, so the topology is computed from the data rather than stored. On this graph the example runs a message-passing sum — every node i collects, from each neighbor j inside the cutoff, the neighbor's value x[j] weighted by the distance between them. It then asks for the gradient of the output with respect to the positions and the values: a vector–Jacobian product (VJP), the backward pass of the same computation.

\[ \text{out}_i \;=\; \sum_{e\,=\,(j \to i),\; \lVert p_j - p_i \rVert \,\le\, r} \lVert p_j - p_i \rVert \cdot x_j \]

The example uses Torch tensors and torch.autograd.grad. Geometry fields are differentiable for the realized neighbor set; the discrete decision to add or remove an edge at the cutoff is not differentiated.

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

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


positions = torch.tensor(
    [[0.0, 0.0], [0.3, 0.0], [0.8, 0.0]], requires_grad=True)  # (N, 2)
source = torch.tensor([2.0, 3.0, 5.0], requires_grad=True)  # (N,)
graph = tg.Graph.radius(positions, cutoff=0.6)

kernel = DistanceWeightedSum()
output = kernel(graph=graph, src={"x": source}, dst={"x": source})  # (N,)
# Edges (dist ≤ 0.6): 1→0 (0.3), 0→1 (0.3), 2→1 (0.5), 1→2 (0.5)
# out = [0.3·3, 0.3·2 + 0.5·5, 0.5·3]

position_grad, source_grad = torch.autograd.grad(
    output,
    (positions, source),
    grad_outputs=torch.tensor([1.0, 2.0, 3.0]),
)

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

"""dynamic radius MessagePassing with automatic geometry VJP."""

import torch
import tiga as tg


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

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


positions = torch.tensor(
    [[0.0, 0.0], [0.3, 0.0], [0.8, 0.0]], requires_grad=True)  # (N, 2)
source = torch.tensor([2.0, 3.0, 5.0], requires_grad=True)  # (N,)
graph = tg.Graph.radius(positions, cutoff=0.6)

kernel = DistanceWeightedSum()
output = kernel(graph=graph, src={"x": source}, dst={"x": source})  # (N,)
# Edges (dist ≤ 0.6): 1→0 (0.3), 0→1 (0.3), 2→1 (0.5), 1→2 (0.5)
# out = [0.3·3, 0.3·2 + 0.5·5, 0.5·3]

position_grad, source_grad = torch.autograd.grad(
    output,
    (positions, source),
    grad_outputs=torch.tensor([1.0, 2.0, 3.0]),
)

# Inspect the selected plan: kernel.explain()

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

CPUCUDA

### DistanceWeightedSum
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: materialize-fixed-radius-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

### DistanceWeightedSum
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: materialize-fixed-radius-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

Exact kNN feeding MessagePassing

What it is. k-nearest neighbors (kNN): given N points, find for each point the k closest other points (here by squared Euclidean distance, self excluded). The example turns each point's k neighbors into graph edges and runs a message-passing sum over them: every point accumulates its neighbors' values x[j] multiplied by an edge weight (all ones here, so the result is a plain neighbor sum).

\[ \text{out}_i \;=\; \sum_{j\,\in\,\mathrm{knn}(i,\,k)} w_{(j \to i)} \cdot x_j \]

Exact kNN stays procedural: the current position snapshot is rebuilt lazily and its fixed-k compressed sparse row (CSR) ABI is rebound to one generated MessagePassing TTIR consumer. Selection and consume are separate contracts, fused into a single launch — no CSR row pointer or selected-column tensor is materialized.

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

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


k = 16
graph = tg.Graph.knn(positions, k)
kernel = NeighborSum()
result = kernel(  # (N,)
    graph=graph, src={"x": values}, dst={},
    edge={"weight": torch.ones(1024 * k, device="cuda")})  # (N·k,)

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

"""Exact dynamic kNN with one compiler-generated, snapshot-rebound consumer."""

import torch

import tiga as tg

positions = torch.rand(1024, 3, device="cuda")                          # (N, 3)
values = torch.rand(1024, device="cuda")                                # (N,)


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

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


k = 16
graph = tg.Graph.knn(positions, k)
kernel = NeighborSum()
result = kernel(  # (N,)
    graph=graph, src={"x": values}, dst={},
    edge={"weight": torch.ones(1024 * k, device="cuda")})  # (N·k,)

# Inspect: kernel.explain(), kernel.ir("kernel")

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

CUDA

### NeighborSumTorchExecutor
backend: cuda
provider: triton-nvidia@3.6.0
lowering: gf-kernel-to-ttir-ranked-select-consume
passes: gf-verify-domain, gf-lower-domain-to-iter, gf-lower-iter-to-kernel, capture-ranked-relation, select-ranked-pairs-coordinate-hierarchy, select-candidate-tile, emit-local-stable-topk, emit-hierarchical-topk-merge, fuse-selected-edge-consumer, gf-kernel-to-ttir, provider-compile-serialized-ttir
provider cache key: triton-nvidia | cuda:nvidia | 3.6.0 | 2.11.0+cu128 | af81e84448f193cc | cuda:120:warp32
remark: [planning] exact candidate ranking and edge consumption share one launch
remark: [planning] candidate tiles retain a power-of-two padded stable key state
remark: [planning] no CSR row pointer or selected column tensor is materialized
variant cache: hits=0, misses=1