fix(zygote ext): handle NamedTuple cotangents in vofa_u_adjoint (EnsembleSolution end-time grads)

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "4.3.2"
+version = "4.3.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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))...),)

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