Repository navigation
Conversation
Fixes #710. - GNNGraphs accepts traced index arrays in a `GNNGraph`, and a new Reactant extension makes `degree` work inside a trace. - A new GNNlib Reactant extension makes `propagate` fall back to gather/scatter for the `copy_xj`, `e_mul_xj` and `w_mul_xj` fast paths, as on CUDA. - GNNlib's node and edge reductions and softmaxes pass explicit output sizes to their scatters, which Reactant needs for traced indices. This also fixes `reduce_nodes`/`reduce_edges` returning too few columns for batches that end in graphs with no nodes or edges. GNNLux gets tests for the forward pass under Reactant and for Enzyme gradients and training inside a compiled program. GNNlib goes to 1.5.0-DEV and requires GNNGraphs 1.6. The root `docs` project leaves the workspace: MultiDocumenter pins JSON 0.21, which holds Reactant at 0.2.217, a version that does not build on Julia 1.13.
Train the tutorials with `Training.single_train_step!(AutoEnzyme(), ...)` on a `reactant_device()`, as Lux recommends, instead of Zygote, and regenerate their pages.
simsurace
marked this pull request as draft
November 5, 2026 16:40
…torial pad_graph pads a graph to fixed numbers of nodes, edges and graphs, so Reactant can compile a model or training step once and reuse it for every mini-batch. The graph classification tutorial goes back to master's mini-batch training, padding each batch and masking the padding graphs out of the loss.
graph_sizes holds the number of nodes in each graph of a batch. It duplicates graph_indicator, but its length is the number of graphs, which makes that number part of the array shapes a compiled Reactant function checks. A compiled model called on a batch with a different number of graphs now raises an error instead of silently returning wrong results. The positional constructor derives graph_sizes from graph_indicator, and transformations that keep the graph indicator pass graph_sizes through.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #710. Follows up on #655.
GNNLux models can now be compiled with Reactant and differentiated with Enzyme, which is the standard Lux workflow.
reactant_device().propagate(copy_xj,e_mul_xj,w_mul_xj) can't be traced, so it falls back to gather/scatter. Scatter sizes are now passed explicitly.BitVectormasks to the device.Upstream context: EnzymeAD/Reactant.jl#2856 and EnzymeAD/Reactant.jl#2875. The losses select the training nodes with
findall(mask)because of a Reactant bug with Boolean indexing (EnzymeAD/Reactant.jl#3087, fixed by EnzymeAD/Reactant.jl#3700). Once that is released, the tutorials can index with the masks directly.Version bumps: GNNGraphs 1.7.0-DEV and GNNlib 1.5.0-DEV, with GNNlib requiring GNNGraphs ≥ 1.7.