From c9fd4a76f18ee7de0d743f02953c12861991cfb8 Mon Sep 17 00:00:00 2001
From: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Date: Mon, 29 Jun 2026 12:21:44 -0400
Subject: [PATCH] fix(zygote ext): pass through NamedTuple cotangents in
 vofa_u_adjoint

`vofa_u_adjoint` rewrapped the per-element cotangents of an
`AbstractVectorOfArray`/`AbstractDiffEqArray` in a `VectorOfArray`/`DiffEqArray`.
For an `EnsembleSolution` whose elements are solution objects, differentiating
only a scalar field (e.g. `sol.t[end]`) makes Zygote produce structural
`NamedTuple` cotangents per trajectory. These have no `ndims`, so the
reconstruction threw `MethodError: no method matching ndims(::NamedTuple)`.

Detect non-array cotangents and pass them through as a plain vector instead.
This extends the earlier `ZeroTangent`/`AbstractZero` handling (#606) to the
structural-tangent case.

Note: this is a local robustness fix only; it does NOT fix
SciML/DifferentialEquations.jl#1149. The cotangent then fails downstream in
SciMLBase's `EnsembleSolution_adjoint`, and SciMLSensitivity never consumes a
`t` cotangent (verified: single-solve `sol.t[end]` loss silently returns zero
gradient under GaussAdjoint while ForwardDiff gives the true nonzero one).
See the PR description for full verification details.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
---
 Project.toml                        |  2 +-
 ext/RecursiveArrayToolsZygoteExt.jl | 11 +++++++++++
 test/AD/adjoints.jl                 | 24 ++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

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 <accounts@chrisrackauckas.com>"]
-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