diff --git a/Project.toml b/Project.toml index 4a5ec21b..ea88fc9e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "RecursiveArrayTools" uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" authors = ["Chris Rackauckas "] -version = "4.3.2" +version = "4.3.3" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index 5dc9032f..634bff0f 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -94,6 +94,7 @@ function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractVectorOfArray) (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i end + _vofa_cotangent_array(m) || return m return VectorOfArray(m) end @@ -102,9 +103,19 @@ function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractDiffEqArray) (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i end + _vofa_cotangent_array(m) || return m return DiffEqArray(m, A.t) end +# Whether the per-element cotangents can be rewrapped in a `VectorOfArray`/`DiffEqArray`. +# When they are structural tangents instead (e.g. `NamedTuple`s for the solution +# objects in an `EnsembleSolution`, produced when only a scalar field such as +# `sol.t[end]` is differentiated) they have no `ndims`, so the cotangent is passed +# through as a plain vector rather than erroring. (DifferentialEquations.jl#1149) +function _vofa_cotangent_array(m) + return all(x -> x isa Union{AbstractArray, AbstractVectorOfArray}, m) +end + @adjoint function literal_getproperty(A::ArrayPartition, ::Val{:x}) function literal_ArrayPartition_x_adjoint(d) (ArrayPartition((isnothing(d[i]) || d[i] isa AbstractZero ? zero(A.x[i]) : d[i] for i in 1:length(d))...),) diff --git a/test/AD/adjoints.jl b/test/AD/adjoints.jl index 4bb1a644..2c48bcac 100644 --- a/test/AD/adjoints.jl +++ b/test/AD/adjoints.jl @@ -118,3 +118,27 @@ let ext = Base.get_extension(RecursiveArrayTools, :RecursiveArrayToolsZygoteExt) @test g_dea.u == expected @test g_dea.t == dea.t end + +# Structural `NamedTuple` cotangents (e.g. the per-trajectory solution cotangents +# of an `EnsembleSolution` when only a scalar field such as `sol.t[end]` is +# differentiated) cannot be wrapped in a `VectorOfArray`/`DiffEqArray`; they must +# pass through as a plain vector instead of erroring on `ndims(::NamedTuple)`. +# (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 = [(u = nothing, t = Float64[0, 0, i]) for i in 1:3] + + g_voa = ext.vofa_u_adjoint(d, voa) + @test g_voa == d + + g_dea = ext.vofa_u_adjoint(d, dea) + @test g_dea == d + + # `nothing`/`AbstractZero` slices still drop out even when mixed with structural tangents. + d_mixed = Any[ + (u = nothing, t = Float64[0, 0, 1]), nothing, + ChainRulesCore.ZeroTangent(), + ] + @test ext.vofa_u_adjoint(d_mixed, voa)[1] == d_mixed[1] +end