fix: make (multi)basis work on CuArray

.buildkite/pipeline.yml

@@ -1,13 +1,11 @@
 steps:
   - label: "DI GPU tests"
-    if: |
-      !build.pull_request.draft &&
-        build.pull_request.labels includes "gpu"
+    if: build.pull_request.labels includes "gpu"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"
     command: |
-      julia ./DifferentiationInterface/test/GPU/CUDA/simple.jl
+      julia ./DifferentiationInterface/test/GPU/CUDA/main.jl
     agents:
       queue: "juliagpu"
       cuda: "*"

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.7.1]
+
+### Fixed
+
+- Make basis work for `CuArray` ([#810])
+
 ## [0.7.0]
 
 ### Changed
@@ -27,11 +33,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allocate Enzyme shadow memory during preparation ([#782])
 
-[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...main
+[unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.1...main
+[0.7.1]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.0...DifferentiationInterface-v0.7.1
 [0.7.0]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.54...DifferentiationInterface-v0.7.0
 [0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
+[#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810
 [#799]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/799
 [#795]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/795
 [#790]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/790

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -41,7 +41,9 @@ DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
 DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
-DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
+DifferentiationInterfacePolyesterForwardDiffExt = [
+    "PolyesterForwardDiff", "ForwardDiff", "DiffResults"
+]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"
 DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
@@ -121,4 +123,21 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]
+test = [
+    "ADTypes",
+    "Aqua",
+    "ComponentArrays",
+    "DataFrames",
+    "ExplicitImports",
+    "JET",
+    "JLArrays",
+    "JuliaFormatter",
+    "Pkg",
+    "Random",
+    "SparseArrays",
+    "SparseConnectivityTracer",
+    "SparseMatrixColorings",
+    "StableRNGs",
+    "StaticArrays",
+    "Test",
+]

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)

DifferentiationInterface/test/Core/Internals/basis.jl

@@ -1,4 +1,4 @@
-using DifferentiationInterface: basis
+using DifferentiationInterface: basis, multibasis
 using LinearAlgebra
 using StaticArrays, JLArrays
 using Test
@@ -8,6 +8,9 @@ using Test
     @test basis(rand(3), 2) isa Vector
     @test basis(rand(3), 2) == b_ref
     @test basis(jl(rand(3)), 2) isa JLArray
+    @test Array(basis(jl(rand(3)), 2)) == [0, 1, 0]
+    @test multibasis(jl(rand(3)), [1, 2]) isa JLArray
+    @test Array(multibasis(jl(rand(3)), [1, 2])) == [1, 1, 0]
     @test all(basis(jl(rand(3)), 2) .== b_ref)
     @test basis(@SVector(rand(3)), 2) isa SVector
     @test basis(@SVector(rand(3)), 2) == b_ref

DifferentiationInterface/test/GPU/CUDA/main.jl

@@ -0,0 +1,9 @@
+@info "Testing on CUDA"
+using Pkg
+Pkg.add("CUDA")
+Pkg.develop(PackageSpec(; path="./DifferentiationInterface"))
+using Test
+
+@testset verbose = true "Simple" begin
+    include("simple.jl")
+end

DifferentiationInterface/test/GPU/CUDA/simple.jl

@@ -1,5 +1,24 @@
-@info "Testing on CUDA"
-using Pkg
-Pkg.add("CUDA")
 using CUDA
+using DifferentiationInterface
+import DifferentiationInterface as DI
+using LinearAlgebra
+using Test
+
 CUDA.versioninfo()
+
+@testset "Basis" begin
+    x = CuVector(rand(Float32, 3))
+    b = DI.basis(x, 2)
+    @test Array(b) == [0, 1, 0]
+
+    X = CuMatrix(rand(Float32, 2, 2))
+    B = DI.multibasis(X, [2, 3])
+    @test Array(B) == [0 1; 1 0]
+end
+
+@testset "Jacobian" begin
+    x = CuVector(rand(Float32, 3))
+    backend = DI.AutoSimpleFiniteDiff()
+    J = jacobian(identity, backend, x)
+    @test (J .!= 0) == I
+end