Skip to content

Commit c547a94

Browse files
Place Mooncake in [weakdeps] and add regression test
The cherry-pick from the v3-backport branch left the Mooncake entry under [deps] rather than [weakdeps] because the surrounding lines in master's Project.toml differ from v3-backport's. Move Mooncake to the weakdeps block so the extension loads via the normal weakdep trigger instead of as a hard dependency, and add Mooncake to [extras] / [targets.test] so the extension is actually exercised in CI. Add test/mooncake.jl with a direct unit test for the new `Mooncake.increment_and_get_rdata!(::FData{@NamedTuple{x::T}}, ::NoRData, ::ArrayPartition{P, T})` dispatch: constructs a matching FData and ArrayPartition, calls `increment_and_get_rdata!`, and checks that (a) the in-place accumulation on each inner-array leaf is correct and (b) the method returns `NoRData()`. Also exercises a three-way Float32 ArrayPartition to cover a different eltype and arity. Register the testset in runtests.jl under the Core group. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 59a96eb commit c547a94

File tree

3 files changed

+45
-4
lines changed

3 files changed

+45
-4
lines changed

Project.toml

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "RecursiveArrayTools"
22
uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
3-
version = "4.0.1"
43
authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
4+
version = "4.0.1"
55

66
[deps]
77
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -38,8 +38,8 @@ RecursiveArrayToolsFastBroadcastPolyesterExt = ["FastBroadcast", "Polyester"]
3838
RecursiveArrayToolsForwardDiffExt = "ForwardDiff"
3939
RecursiveArrayToolsKernelAbstractionsExt = "KernelAbstractions"
4040
RecursiveArrayToolsMeasurementsExt = "Measurements"
41-
RecursiveArrayToolsMooncakeExt = "Mooncake"
4241
RecursiveArrayToolsMonteCarloMeasurementsExt = "MonteCarloMeasurements"
42+
RecursiveArrayToolsMooncakeExt = "Mooncake"
4343
RecursiveArrayToolsReverseDiffExt = ["ReverseDiff", "Zygote"]
4444
RecursiveArrayToolsSparseArraysExt = ["SparseArrays"]
4545
RecursiveArrayToolsStatisticsExt = "Statistics"
@@ -60,8 +60,8 @@ GPUArraysCore = "0.2"
6060
KernelAbstractions = "0.9.36"
6161
LinearAlgebra = "1.10"
6262
Measurements = "2.11"
63-
Mooncake = "0.5"
6463
MonteCarloMeasurements = "1.2"
64+
Mooncake = "0.5"
6565
NLsolve = "4.5"
6666
Pkg = "1"
6767
Polyester = "0.7.16"
@@ -89,6 +89,7 @@ FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
8989
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
9090
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
9191
Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
92+
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
9293
MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
9394
NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
9495
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -105,4 +106,4 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
105106
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
106107

107108
[targets]
108-
test = ["Aqua", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "Measurements", "NLsolve", "Pkg", "Polyester", "Random", "SafeTestsets", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful", "Zygote"]
109+
test = ["Aqua", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "Measurements", "Mooncake", "NLsolve", "Pkg", "Polyester", "Random", "SafeTestsets", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful", "Zygote"]

test/mooncake.jl

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
using RecursiveArrayTools, Mooncake, Test
2+
3+
# Regression test for the `RecursiveArrayToolsMooncakeExt` dispatch that
4+
# lets Mooncake's `@from_chainrules`/`@from_rrule` accumulator handle an
5+
# `ArrayPartition` cotangent returned by an upstream ChainRule (e.g.
6+
# SciMLSensitivity's `_concrete_solve_adjoint` for a `SecondOrderODEProblem`).
7+
# Without the extension, the call below fell through to Mooncake's generic
8+
# error path:
9+
#
10+
# ArgumentError: The fdata type Mooncake.FData{@NamedTuple{x::Tuple{Vector{Float64}, Vector{Float64}}}},
11+
# rdata type Mooncake.NoRData, and tangent type
12+
# RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}
13+
# combination is not supported with @from_chainrules or @from_rrule.
14+
15+
@testset "ArrayPartition increment_and_get_rdata!" begin
16+
@test Base.get_extension(RecursiveArrayTools, :RecursiveArrayToolsMooncakeExt) !==
17+
nothing
18+
19+
# Tangent produced by an upstream ChainRule.
20+
t = ArrayPartition([1.0, 2.0], [3.0, 4.0])
21+
# Pre-existing FData that the method should accumulate into in place.
22+
f = Mooncake.FData((x = ([10.0, 20.0], [30.0, 40.0]),))
23+
24+
r = Mooncake.increment_and_get_rdata!(f, Mooncake.NoRData(), t)
25+
26+
@test r === Mooncake.NoRData()
27+
@test f.data.x[1] == [11.0, 22.0]
28+
@test f.data.x[2] == [33.0, 44.0]
29+
30+
# Three-way partition with Float32 leaves — exercises the inner
31+
# per-leaf dispatch on a different eltype and arity.
32+
t32 = ArrayPartition(Float32[1, 2], Float32[3, 4, 5], Float32[6])
33+
f32 = Mooncake.FData((x = (Float32[10, 20], Float32[30, 40, 50], Float32[60]),))
34+
r32 = Mooncake.increment_and_get_rdata!(f32, Mooncake.NoRData(), t32)
35+
@test r32 === Mooncake.NoRData()
36+
@test f32.data.x[1] == Float32[11, 22]
37+
@test f32.data.x[2] == Float32[33, 44, 55]
38+
@test f32.data.x[3] == Float32[66]
39+
end

test/runtests.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ end
4646
@time @safetestset "StaticArrays Tests" include("copy_static_array_test.jl")
4747
@time @safetestset "Linear Algebra Tests" include("linalg.jl")
4848
@time @safetestset "Adjoint Tests" include("adjoints.jl")
49+
@time @safetestset "Mooncake Tests" include("mooncake.jl")
4950
@time @safetestset "Measurement Tests" include("measurements.jl")
5051
end
5152

0 commit comments

Comments
 (0)