fix: make (multi)basis work on CuArray

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

@@ -1,42 +1,21 @@
 module DifferentiationInterfaceGPUArraysCoreExt
 
 import DifferentiationInterface as DI
-using GPUArraysCore: AbstractGPUArray
+using GPUArraysCore: @allowscalar, AbstractGPUArray
 
-"""
-    OneElement
-
-Efficient storage for a one-hot array, aka an array in the standard Euclidean basis.
-"""
-struct OneElement{I,N,T,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
-    ind::I
-    val::T
-    a::A
-
-    function OneElement(ind::Integer, val::T, a::A) where {N,T,A<:AbstractArray{T,N}}
-        right_ind = eachindex(a)[ind]
-        return new{typeof(right_ind),N,T,A}(right_ind, val, a)
-    end
-
-    function OneElement(
-        ind::CartesianIndex{N}, val::T, a::A
-    ) where {N,T,A<:AbstractArray{T,N}}
-        linear_ind = LinearIndices(a)[ind]
-        right_ind = eachindex(a)[linear_ind]
-        return new{typeof(right_ind),N,T,A}(right_ind, val, a)
-    end
-end
-
-Base.size(oe::OneElement) = size(oe.a)
-Base.IndexStyle(oe::OneElement) = Base.IndexStyle(oe.a)
-
-function Base.getindex(oe::OneElement{<:Integer}, ind::Integer)
-    return ifelse(ind == oe.ind, oe.val, zero(eltype(oe.a)))
+function DI.basis(a::AbstractGPUArray{T}, i) where {T}
+    b = similar(a)
+    fill!(b, zero(T))
+    @allowscalar b[i] = one(T)
+    return b
 end
 
-function DI.basis(a::AbstractGPUArray{T}, i) where {T}
-    b = zero(a)
-    b .+= OneElement(i, one(T), a)
+function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
+    b = similar(a)
+    fill!(b, zero(T))
+    for i in inds
+        @allowscalar b[i] = one(T)
+    end
     return b
 end
 

DifferentiationInterface/src/utils/basis.jl

@@ -18,9 +18,6 @@ end
     multibasis(a::AbstractArray, inds)
 
 Construct the sum of the `i`-th standard basis arrays in the vector space of `a` for all `i ∈ inds`.
-
-!!! warning
-    Does not work on GPU, since this is intended for sparse autodiff and SparseMatrixColorings.jl doesn't work on GPUs either.
 """
 function multibasis(a::AbstractArray{T}, inds) where {T}
     b = similar(a)