diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index 5dc9032f..1e0a6909 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -90,6 +90,11 @@ Zygote.@adjoint function Zygote.literal_getproperty(A::RecursiveArrayTools.Abstr end function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractVectorOfArray) + # Structural tangents (e.g. per-solution NamedTuples for an + # EnsembleSolution's `u`) cannot be stored in a VectorOfArray, so the + # cotangent of `A` is the `u`-field tangent itself. + all(d_i -> isnothing(d_i) || d_i isa Union{AbstractZero, AbstractArray, Number}, d) || + return (; u = collect(d)) m = map(enumerate(d)) do (idx, d_i) (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i @@ -98,6 +103,8 @@ function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractVectorOfArray) end function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractDiffEqArray) + all(d_i -> isnothing(d_i) || d_i isa Union{AbstractZero, AbstractArray, Number}, d) || + return (; u = collect(d)) m = map(enumerate(d)) do (idx, d_i) (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i diff --git a/test/AD/adjoints.jl b/test/AD/adjoints.jl index 4bb1a644..93e6d714 100644 --- a/test/AD/adjoints.jl +++ b/test/AD/adjoints.jl @@ -118,3 +118,24 @@ let ext = Base.get_extension(RecursiveArrayTools, :RecursiveArrayToolsZygoteExt) @test g_dea.u == expected @test g_dea.t == dea.t end + +# Structural per-element cotangents (e.g. the per-solution NamedTuple tangents +# produced when differentiating `ensemble_sol.u`) cannot be stored in a +# VectorOfArray; the `u`-field tangent is returned instead +# (SciML/DifferentialEquations.jl#1149). +let ext = Base.get_extension(RecursiveArrayTools, :RecursiveArrayToolsZygoteExt) + voa = VectorOfArray([Float64[i, i, i] for i in 1:3]) + dea = DiffEqArray([Float64[i, i, i] for i in 1:3], 1:3) + d = Any[ + (; u = Float64[1, 1, 1], t = 0.5), nothing, + (; u = Float64[3, 3, 3], t = 0.25), + ] + + g_voa = ext.vofa_u_adjoint(d, voa) + @test g_voa isa NamedTuple + @test g_voa.u == collect(d) + + g_dea = ext.vofa_u_adjoint(d, dea) + @test g_dea isa NamedTuple + @test g_dea.u == collect(d) +end