Sparse aggregation feasibility¶

This record contains revision-bound implementation evidence. Read Sparse graph topology first for the stable concepts.

Decision¶

At tinygrad revision b1a7229, a tinymesh-owned destination-CSR kernel passes the identity-message correctness, gradient, and sparse-structure gates on CPU and Metal. Forward computes A @ X; its custom gradient runs the same kernel over transpose CSR to compute A^T @ dY. Neither direction uses atomics or materializes [E, H] state.

This resolves the implementation boundary without adding a tinygrad primitive. The kernel now lives behind experimental Graph.sum; the benchmark remains in experiments/. The backend relies on the alpha Tensor.custom_kernel surface and a data-dependent UOp.loop. At recorded revision 0bb36c9, tinygrad's default kernel optimization rewrites this loop into invalid UOps and fails type verification, so the kernel pins opts_to_apply=(). It covers fixed-topology first-order sum aggregation only and serializes each CSR row within one feature lane. Those constraints limit stability even though the narrow operation now has a package owner.

CSR result¶

One immutable topology stores both destination CSR and its transpose:

forward:  destination row -> source columns -> sum source features
backward: source row      -> target columns -> sum target gradients

The transpose makes backward the same operation as forward. Preprocessing sorts each row by column, so a fixed node labeling is deterministic regardless of COO input order. Duplicate edges retain multiplicity; isolated rows return zero. Hand-computed tests cover both directions, repeated sources and destinations, duplicate edges, isolated nodes, empty topology, COO-order invariance, and vertex-permutation equivariance on exact-valued fixtures on CPU and Metal. As with any sequential floating-point reduction, relabeling can change summation order and therefore the final rounding bits.

For balanced degree-eight graphs with feature width H = 32, the two CSR connectivity forms have exactly 2E + 2(N + 1) integer elements. Each direction launches N * H lanes whose row loops make exactly E * H active edge visits:

Doubling N and E therefore doubles useful lane work exactly and grows CSR storage linearly. The current topology also retains two edge-order maps of length E; they preserve scalar edge identity without changing the complexity. This follows from the representation and the single dynamic row loop; it is the complexity proof. tinygrad's counters cannot observe data-dependent loop trips, so wall time is supporting evidence, not the proof.

On a 32 GB Apple M4 MacBook Air running macOS 26.5.2, after five warmups and across 20 samples, the largest cases produced these representative medians:

The hub cases expose the present tradeoff: one high-degree row serializes its edges across only H feature lanes. Preprocessing and compilation are excluded from these kernel medians. tinygrad's kernel option search is disabled as noted above, and repeated balanced timings varied enough that no performance claim should rest on the table alone.

The benchmark does not time weighted or edge-dependent messages. A separate weighted aggregation record proves scalar weighted forward, dX, and dw without materializing [E, H] state. Vector edge messages, batching, changing topology, higher-order gradients, and topology construction cost remain open. Fixed-topology buffers are realized once per device and reused by subsequent calls; preprocessing remains a host-side cost paid by each new topology. Empty graphs use tinygrad's ordinary values * 0 path because a custom kernel cannot receive zero-length buffers. A one-node graph similarly reduces to values * E, avoiding unnecessary traversal and a degenerate Metal loop scope.

Rejected public composition¶

At the earlier tinygrad revision 980748c, the smallest native candidate was:

node state [N, H]
  -> gather source rows         [E, H]
  -> scatter_reduce(sum) by target [N, H]

It passed hand-computed forward values, isolated-node behavior, finite-difference gradients, and vertex-permutation equivariance. It failed the defining sparse invariant: work must scale with edges rather than node-edge pairs.

Evidence¶

For E = 2N and H = 4, doubling both N and E should approximately double the work of an edge-linear implementation. The native candidate approaches a fourfold increase instead:

A fresh pinned-revision Metal rerun retained 24 N * E * H carriers at every measured size. tinygrad may fuse these rather than allocate the full tensors, but fusion does not remove the quadratic work. A synthetic N=30,000, E=60,000, H=32 case therefore implies 57.6 billion logical node-edge-feature lanes before useful message computation.

This matches the pinned implementation: gather uses a one-hot mask and scatter_reduce expands source values and a destination mask before reduction (source). The current upstream scatter-reduce comparisons remain forward-only (tests), although this experiment observes correct sum gradients on its small cases.

tinygrad's atomic embedding backward was not a complete tested alternative. It is a custom UOp kernel limited to CPU and AMD, while Tensor.custom_kernel is marked alpha; its forward lookup still uses one-hot reduction (implementation). This leaves custom UOps as a research candidate rather than an established library path; it does not assign the solution to either repository.

Precursor custom-kernel probe¶

An isolated 2026-07-21 probe at tinygrad f64f96ec made the unresolved boundary narrower. Current source still uses one-hot expansion for public gather and scatter_reduce, while Tensor.custom_kernel remains alpha (indexed operations, custom kernels):

The direct probes used existing UOps through Tensor.custom_kernel; they did not require a new low-level operation. The CSR probe avoided atomics by giving each destination row one writer. The checked-in experiment above completes its first-order gradient with transpose CSR.

At that revision, default kernel optimization failed with KeyError: dtypes.weakint unless options were fixed explicitly, and normal statistics collection raised TypeError: _f() missing 1 required positional argument: 'core_id' for the data-dependent row loop. tinygrad revision b1a7229 supplies the loop form used by the checked-in experiment, but default kernel optimization still fails with the same KeyError. At current revision 0bb36c9, that failure instead reaches type verification and raises RuntimeError: UOp verification failed on an Ops.END or Ops.RESHAPE, depending on the fixture. The candidate therefore still pins opts_to_apply=(). The alpha API, disabled option search, and unmeasured compilation and topology-preprocessing costs remain limitations.

How the PyTorch stack expresses it¶

PyTorch Geometric 2.8.0 uses the same semantic lowering tinymesh needs, but its backend primitives are genuinely indexed:

node state [N, H]
  -> torch.index_select(source)       [E, H]
  -> message                          [E, H]
  -> zeros[N, H].scatter_add(target)  [N, H]

MessagePassing obtains source or target rows with Tensor.index_select (PyG source). Its sum and mean aggregations use PyTorch's in-place scatter_add_ (PyG source). The expanded target index has shape [E, H]; it never introduces an N * E axis. For compatible layers and sparse adjacency inputs, PyG can instead fuse message and aggregation into sparse matrix multiplication (PyG source).

PyTorch owns the corresponding backward operations. The source gradient of scatter_add is a gather (PyTorch source); the input gradient of index_select is a zero tensor followed by index_add_ (PyTorch source). Both directions therefore touch node or edge feature lanes, not every node-edge pair. PyG owns graph semantics and orchestration; PyTorch owns indexed loads, accumulation, autograd, and device kernels.

PyTorch Geometric Temporal adds no universal temporal graph kernel. Its static signal container returns one ordinary PyG Data snapshot at a time while reusing one edge_index (source). The recurrent model then composes PyG operations across time:

• GConvGRU performs six ChebConv calls per time step—two for each GRU gate—and each Chebyshev order adds another sparse propagation (cell, convolution). It remains edge-linear, with work roughly proportional to T * K * E times feature widths and gate constants.
• The ordinary DCRNN implementation is not sparse end to end. Its DConv constructs an [N, N] adjacency with to_dense_adj for degree calculation and reverse-edge construction before calling sparse propagate (source). DCRNN invokes that layer for all three recurrent gates. The separate BatchedDConv path instead computes degrees with scatter_add_ and reverses the edge list directly (source).

The lesson is narrower than “PyG Temporal is sparse”: PyG's primitive path is edge-linear, but every temporal architecture must still be audited for dense normalization, adjacency conversion, batching, and repeated topology work.

Reproduce¶

uv run --locked python -m unittest tests.test_sparse_aggregation
DEV=CPU uv run --locked python -m unittest tests.test_graph
DEV=METAL uv run --locked python -m unittest tests.test_graph
uv run --locked python -m experiments.run sparse_aggregation DEV=CPU
uv run --locked python -m experiments.run csr_aggregation DEV=CPU
uv run --locked python -m experiments.run csr_aggregation DEV=METAL

The CSR experiment satisfies the structural gate for fixed-topology, identity-message, first-order sum aggregation in both directions. More complex messages add edge-local costs and require their own proof. Mean GraphSAGE and unweighted GCN learn from the shared sum output; dedicated graph tests own its transpose-gradient evidence. Neither makes the alpha custom-kernel or dynamic-loop contracts stable. The rejected public gather-and-scatter composition must not be called sparse graph support. A later T-GCN experiment proves that the same fixed Graph can be reused across ordered snapshots.