From f9610ad66f4a068bdec411070c94d8f9a01eed12 Mon Sep 17 00:00:00 2001
From: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Date: Sat, 6 Jun 2026 05:27:42 -0400
Subject: [PATCH] fix: handle ZeroTangent/AbstractZero in
 RecursiveArrayToolsZygoteExt adjoint

The Zygote pullbacks for VectorOfArray/DiffEqArray (vofa_u_adjoint) and
ArrayPartition only special-cased `isnothing(d_i)` for missing per-slice
gradients. A ChainRulesCore.ZeroTangent (AbstractZero) cotangent element
was passed through unchanged, and the subsequent VectorOfArray/DiffEqArray
construction called `size(::ZeroTangent)`, throwing a MethodError. This
broke reverse-mode AD through DiffEqArray(vec, ts) (e.g. DifferenceEquations.jl
test/linear_gradients.jl).

Treat AbstractZero the same as nothing: a structural zero gradient for that
slice. Added a regression test in test/adjoints.jl exercising a ZeroTangent
cotangent element for both VectorOfArray and DiffEqArray.

Regression test run locally on Julia 1.10: PASS.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 ext/RecursiveArrayToolsZygoteExt.jl |  7 ++++---
 test/adjoints.jl                    | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl
index e4ef4c1f..5dc9032f 100644
--- a/ext/RecursiveArrayToolsZygoteExt.jl
+++ b/ext/RecursiveArrayToolsZygoteExt.jl
@@ -4,6 +4,7 @@ using RecursiveArrayTools
 
 using Zygote
 using Zygote: FillArrays, ChainRulesCore, literal_getproperty, @adjoint
+using Zygote.ChainRulesCore: AbstractZero
 
 function ChainRulesCore.rrule(
         T::Type{<:RecursiveArrayTools.GPUArraysCore.AbstractGPUArray},
@@ -90,7 +91,7 @@ end
 
 function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractVectorOfArray)
     m = map(enumerate(d)) do (idx, d_i)
-        isnothing(d_i) && return zero(A.u[idx])
+        (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx])
         d_i
     end
     return VectorOfArray(m)
@@ -98,7 +99,7 @@ end
 
 function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractDiffEqArray)
     m = map(enumerate(d)) do (idx, d_i)
-        isnothing(d_i) && return zero(A.u[idx])
+        (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx])
         d_i
     end
     return DiffEqArray(m, A.t)
@@ -106,7 +107,7 @@ end
 
 @adjoint function literal_getproperty(A::ArrayPartition, ::Val{:x})
     function literal_ArrayPartition_x_adjoint(d)
-        (ArrayPartition((isnothing(d[i]) ? zero(A.x[i]) : d[i] for i in 1:length(d))...),)
+        (ArrayPartition((isnothing(d[i]) || d[i] isa AbstractZero ? zero(A.x[i]) : d[i] for i in 1:length(d))...),)
     end
     A.x, literal_ArrayPartition_x_adjoint
 end
diff --git a/test/adjoints.jl b/test/adjoints.jl
index 4e60c108..4bb1a644 100644
--- a/test/adjoints.jl
+++ b/test/adjoints.jl
@@ -1,4 +1,5 @@
 using RecursiveArrayTools, Zygote, ForwardDiff, ReverseDiff, Test
+using Zygote: ChainRulesCore
 
 function loss(x)
     return sum(abs2, Array(VectorOfArray([x .* i for i in 1:5])))
@@ -99,3 +100,21 @@ end
 let g = ReverseDiff.gradient(x -> sum(VectorOfArray([VectorOfArray([x])])), float.(1:3))
     @test g ≈ ones(3)
 end
+
+# A `ZeroTangent`/`AbstractZero` slice in the cotangent (structural zero gradient)
+# must be treated like an absent gradient instead of erroring on `size(::ZeroTangent)`.
+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[Float64[1, 1, 1], ChainRulesCore.ZeroTangent(), Float64[3, 3, 3]]
+    expected = [Float64[1, 1, 1], Float64[0, 0, 0], Float64[3, 3, 3]]
+
+    g_voa = ext.vofa_u_adjoint(d, voa)
+    @test g_voa isa VectorOfArray
+    @test g_voa.u == expected
+
+    g_dea = ext.vofa_u_adjoint(d, dea)
+    @test g_dea isa DiffEqArray
+    @test g_dea.u == expected
+    @test g_dea.t == dea.t
+end