allow cu(::ParamDependentOp) that returns an Array

Author: marius311

src/field_tuples.jl

@@ -42,7 +42,7 @@ getindex(D::DiagOp{<:FieldTuple}, i::Int, j::Int) = (i==j) ? D.diag[:][i] : diag
 
 ### array interface
 size(f::FieldTuple) = (sum(map(length, f.fs)),)
-copyto!(dest::FT, src::FT) where {FT<:FieldTuple} = (map(copyto!,dest.fs,src.fs); dest)
+copyto!(dest::FieldTuple{B}, src::FieldTuple{B}) where {B} = (map(copyto!,dest.fs,src.fs); dest)
 iterate(ft::FieldTuple, args...) = iterate(ft.fs, args...)
 getindex(f::FieldTuple, i::Union{Int,UnitRange}) = getindex(f.fs, i)
 fill!(ft::FieldTuple, x) = (map(f->fill!(f,x), ft.fs); ft)

src/flat_s0.jl

@@ -58,6 +58,7 @@ function similar(f::F,::Type{T},dims::Dims) where {P,F<:FlatS0{P},T<:Number}
     @assert size(f)==dims "Tried to make a field similar to $F but dims should have been $(size(f)), not $dims."
     basetype(F){P}(similar(firstfield(f),T))
 end
+copyto!(dst::Field{B,S0,P}, src::Field{B,S0,P}) where {B,P} = (copyto!(firstfield(dst),firstfield(src)); dst)
 
 
 ### broadcasting

src/gpu.jl

@@ -38,10 +38,6 @@ BroadcastStyle(::FlatS0Style{F,Array}, ::FlatS0Style{F,CuArray}) where {P,F<:Fla
 # the generic versions of these trigger scalar indexing of CUDA, so provide
 # specialized versions: 
 
-function copyto!(dst::F, src::F) where {F<:CuFlatS0}
-    copyto!(firstfield(dst),firstfield(src))
-    dst
-end
 pinv(D::Diagonal{T,<:CuFlatS0}) where {T} = Diagonal(@. ifelse(isfinite(inv(D.diag)), inv(D.diag), $zero(T)))
 inv(D::Diagonal{T,<:CuFlatS0}) where {T} = any(Array((D.diag.==0)[:])) ? throw(SingularException(-1)) : Diagonal(inv.(D.diag))
 fill!(f::CuFlatS0, x) = (fill!(firstfield(f),x); f)

src/specialops.jl

@@ -218,7 +218,17 @@ function (L::ParamDependentOp)(θ::NamedTuple)
         # filtering out non-dependent parameters disabled until I can find a fix to:
         # https://discourse.julialang.org/t/can-zygote-do-derivatives-w-r-t-keyword-arguments-which-get-captured-in-kwargs/34553/8
         # dependent_θ = filter(((k,_),)->k in L.parameters, pairs(θ))
-        L.recompute_function(;θ...)
+        Lθ = L.recompute_function(;θ...)
+        if Lθ isa typeof(L.op)
+            Lθ
+        else
+            # if L got adapt'ed to CuArray since this op was created,
+            # L.op will be GPU-backed, but depending on how
+            # recompute_function is written, recompute_function may
+            # still return something CPU-backed. in that case, copy it
+            # to GPU here
+            copyto!(similar(L.op), Lθ)
+        end
     else
         L.op
     end