Skip to content

Create the eigensolver start vector on the device of the matrix #722

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CarloLucibello merged 2 commits into from
Nov 1, 2026
Merged

Create the eigensolver start vector on the device of the matrix #722

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f359027
Create the eigensolver start vector on the device of the matrix
simsurace Oct 31, 2026
abf2c46
Link the changelog entry to #722
simsurace Oct 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ Entries link to the pull request that introduced them.
- Added a `Mooncake` package extension, so that Mooncake can differentiate on CUDA through `add_self_loops` and through adjacency-matrix graphs. Float edge weights of adjacency-matrix graphs stay differentiable ([#704]).

**Fixed**
- Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph ([#722]).
- Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]).
- Enzyme can now differentiate through the adjacency-matrix → COO graph conversion: a new internal keyword-free helper `_to_coo_graph` avoids the union-typed keyword handling of the `GNNGraph(g; graph_type)` constructor and of `to_coo`, which Enzyme's type analysis cannot compile. Together with upstream fixes in Enzyme ≥ 0.13.197 this removes the `EnzymeInternalError` crash on `:dense`/`:sparse` graphs ([#703]).

Expand Down Expand Up @@ -238,4 +239,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
[#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704
[#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707
[#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706
[#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722
[FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662
5 changes: 4 additions & 1 deletion GNNGraphs/src/query.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -508,7 +508,10 @@ function _eigmax(A)
KrylovKit.eigsolve(Symmetric(A), x0, 1, :LR)[1][1] # also eigs(A, x0, nev, mode) available
end

_rand_dense_vector(A::AbstractMatrix{T}) where {T} = randn(float(T), size(A, 1))
# On the same device as `A`, otherwise a GPU `A` would be passed to CPU BLAS.
_rand_dense_vector(A::AbstractMatrix{T}) where {T} = randn!(similar(A, float(T), size(A, 1)))
# `similar` would give a sparse vector here.
_rand_dense_vector(A::AbstractSparseMatrix{T}) where {T} = randn(float(T), size(A, 1))

# Eigenvalues for cuarray don't seem to be well supported.
# https://github.com/JuliaGPU/CUDA.jl/issues/154
Expand Down
Loading