From b49b97a83d7b107083c83a59927b0949b8c3639e Mon Sep 17 00:00:00 2001
From: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Date: Sat, 19 Sep 2026 18:49:36 -0400
Subject: [PATCH] Return u-field tangent for structural VectorOfArray
 cotangents
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Zygote's `literal_getproperty` adjoint for `A.u` received per-element
cotangents and always rebuilt a VectorOfArray/DiffEqArray from them.
When the elements are structural tangents — e.g. the per-solution
NamedTuples produced when differentiating `ensemble_sol.u` — the
VectorOfArray constructor fails on `ndims(::NamedTuple)`, breaking
ensemble differentiation
(SciML/DifferentialEquations.jl#1149).

Detect non-numeric structural elements (anything that is not nothing,
an AbstractZero, an AbstractArray or a Number) and return the
`u`-field tangent `(; u = collect(d))` instead. Numeric cotangents keep
the existing VectorOfArray/DiffEqArray behavior.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Agent-Harness: Devin CLI 3000.10.31
Agent-Model: SWE-2 Max
Agent-Session: local transcript /home/crackauc/.local/share/devin/cli/summaries/history_d4203d25c3424828.md
---
 ext/RecursiveArrayToolsZygoteExt.jl |  7 +++++++
 test/AD/adjoints.jl                 | 21 +++++++++++++++++++++
 2 files changed, 28 insertions(+)

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