From b3a1b300608cfa48b6dd15d507837ec047f93698 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <katharine.hyatt@ugent.be>
Date: Tue, 18 Aug 2026 13:21:48 +0200
Subject: [PATCH] Add GPU-safe index selection for eig/eigh pullbacks

---
 .../MatrixAlgebraKitCUDAExt.jl                |  8 -------
 src/common/pullbacks.jl                       | 22 +++++++++++++++++++
 src/pullbacks/eig.jl                          |  9 ++++----
 src/pullbacks/eigh.jl                         |  9 ++++----
 4 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl
index 9fa58750f..704e92e7a 100644
--- a/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl
+++ b/ext/MatrixAlgebraKitCUDAExt/MatrixAlgebraKitCUDAExt.jl
@@ -207,14 +207,6 @@ function svd_pullback!(ΔA::AnyCuMatrix, A, USVᴴ, ΔUSVᴴ, ind::AnyCuVector;
     return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, collect(ind); kwargs...)
 end
 
-function eigh_pullback!(ΔA::AnyCuMatrix, A, DV, ΔDV, ind::AnyCuVector; kwargs...)
-    return eigh_pullback!(ΔA, A, DV, ΔDV, collect(ind); kwargs...)
-end
-
-function eig_pullback!(ΔA::AnyCuMatrix, A, DV, ΔDV, ind::AnyCuVector; kwargs...)
-    return eig_pullback!(ΔA, A, DV, ΔDV, collect(ind); kwargs...)
-end
-
 # have to override this as methods are missing in GPUArrays for the various
 # views of Diagonal of ΔA
 function svd_pushforward!(
diff --git a/src/common/pullbacks.jl b/src/common/pullbacks.jl
index 4fe853cdd..053db5335 100644
--- a/src/common/pullbacks.jl
+++ b/src/common/pullbacks.jl
@@ -13,3 +13,25 @@ iszerotangent(::Nothing) = true
 
 # fallback
 _sylvester(A, B, C) = LinearAlgebra.sylvester(A, B, C)
+
+"""
+    select_indices(r::AbstractRange, ind)
+
+Compute `r[ind]` without iterating over `ind`, so that this also works for an `ind` that
+lives on a device.
+"""
+select_indices(r::AbstractRange, ind) = r[ind]
+select_indices(r::AbstractRange, ind::AbstractRange{<:Integer}) = r[ind]
+function select_indices(r::AbstractRange, ind::AbstractVector{<:Integer})
+    checkbounds(r, ind)
+    return first(r) .+ step(r) .* (ind .- 1)
+end
+
+"""
+    is_leading_index(ind, p::Int)
+
+Check whether `ind` selects the first `p` values in order, i.e. whether `ind == 1:p`, without
+iterating over `ind`, so that this also works for an `ind` that lives on a device.
+"""
+is_leading_index(ind::AbstractRange, p::Int) = ind == 1:p
+is_leading_index(ind::AbstractVector, p::Int) = length(ind) == p && all(ind .== 1:p)
diff --git a/src/pullbacks/eig.jl b/src/pullbacks/eig.jl
index 5ec817bf6..ad2fa2464 100755
--- a/src/pullbacks/eig.jl
+++ b/src/pullbacks/eig.jl
@@ -5,12 +5,12 @@ function check_and_prepare_eig_cotangents(
     )
 
     n, p = size(V)
-    indD = axes(D, 1)[ind]
-    indV = axes(V, 2)[ind]
+    indD = select_indices(axes(D, 1), ind)
+    indV = select_indices(axes(V, 2), ind)
     if !iszerotangent(ΔV)
         n == size(ΔV, 1) || throw(DimensionMismatch())
         length(indV) == size(ΔV, 2) || throw(DimensionMismatch())
-        if indV == 1:p
+        if is_leading_index(indV, p)
             ΔV₁ = copy(ΔV)
         else
             ΔV₁ = zero(V)
@@ -41,8 +41,7 @@ function check_and_prepare_eig_cotangents(
     if !iszerotangent(ΔDmat)
         ΔD = diagview(ΔDmat)
         length(indD) == length(ΔD) || throw(DimensionMismatch())
-        # needed to avoid GPUCompiler errors
-        VᴴAΔV[diagind(VᴴAΔV)[indD]] .+= ΔD
+        VᴴAΔV[select_indices(diagind(VᴴAΔV), indD)] .+= ΔD
     else
         ΔD = nothing
     end
diff --git a/src/pullbacks/eigh.jl b/src/pullbacks/eigh.jl
index 20ac7fe5d..cc87ea484 100755
--- a/src/pullbacks/eigh.jl
+++ b/src/pullbacks/eigh.jl
@@ -5,12 +5,12 @@ function check_and_prepare_eigh_cotangents(
     )
 
     n, p = size(V)
-    indD = axes(D, 1)[ind]
-    indV = axes(V, 2)[ind]
+    indD = select_indices(axes(D, 1), ind)
+    indV = select_indices(axes(V, 2), ind)
     if !iszerotangent(ΔV)
         n == size(ΔV, 1) || throw(DimensionMismatch())
         length(indV) == size(ΔV, 2) || throw(DimensionMismatch())
-        if indV == 1:p
+        if is_leading_index(indV, p)
             ΔV₁ = copy(ΔV)
         else
             ΔV₁ = zero(V)
@@ -42,8 +42,7 @@ function check_and_prepare_eigh_cotangents(
     if !iszerotangent(ΔDmat)
         ΔD = diagview(ΔDmat)
         length(indD) == length(ΔD) || throw(DimensionMismatch())
-        # needed to avoid GPUCompiler errors
-        VᴴAΔV[diagind(VᴴAΔV)[indD]] .+= real.(ΔD)
+        VᴴAΔV[select_indices(diagind(VᴴAΔV), indD)] .+= real.(ΔD)
     else
         ΔD = nothing
     end