Fix JLArrays until KA 0.10

Author: christiangnrd

lib/JLArrays/src/JLArrays.jl

@@ -347,6 +347,43 @@ function GPUArrays.mapreducedim!(f, op, R::AnyJLArray, A::Union{AbstractArray,Br
     R
 end
 
+## Base interface
+
+Base._accumulate!(op, output::AnyJLArray, input::AnyJLVector, dims::Nothing, init::Nothing) =
+    accumulate!(op, typed_data(output), typed_data(input); dims=1)
+
+Base._accumulate!(op, output::AnyJLArray, input::AnyJLArray, dims::Integer, init::Nothing) =
+    accumulate!(op, typed_data(output), typed_data(input); dims)
+
+Base._accumulate!(op, output::AnyJLArray, input::AnyJLVector, dims::Nothing, init::Some) =
+    accumulate!(op, typed_data(output), typed_data(input); dims=1, init=something(init))
+
+Base._accumulate!(op, output::AnyJLArray, input::AnyJLArray, dims::Integer, init::Some) =
+    accumulate!(op, typed_data(output), typed_data(input); dims, init=something(init))
+
+Base.accumulate_pairwise!(op, result::AnyJLVector, v::AnyJLVector) = accumulate!(op, result, v)
+
+# default behavior unless dims are specified by the user
+function Base.accumulate(op, A::AnyJLArray;
+                         dims::Union{Nothing,Integer}=nothing, kw...)
+    nt = values(kw)
+    if dims === nothing && !(A isa AbstractVector)
+        # This branch takes care of the cases not handled by `_accumulate!`.
+        return reshape(accumulate(op, typed_data(A)[:]; kw...), size(A))
+    end
+    if isempty(kw)
+        out = similar(A, Base.promote_op(op, eltype(A), eltype(A)))
+        init = AK.neutral_element(op, eltype(out))
+    elseif keys(nt) === (:init,)
+        out = similar(A, Base.promote_op(op, typeof(nt.init), eltype(A)))
+        init = nt.init
+    else
+        throw(ArgumentError("accumulate does not support the keyword arguments $(setdiff(keys(nt), (:init,)))"))
+    end
+    accumulate!(op, typed_data(out), typed_data(A); dims, init)
+end
+
+
 ## KernelAbstractions interface
 
 KernelAbstractions.get_backend(a::JLA) where JLA <: JLArray = JLBackend()

test/Project.toml

@@ -12,3 +12,6 @@ REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+JLArrays = {path = "/Users/christian/.julia/dev/GPUArrays/lib/JLArrays"}