feat: recursive similar for caches

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -89,7 +89,7 @@ function _translate(
 end
 function _translate(::Type{D}, c::DI.Cache) where {D<:Dual}
     c0 = DI.unwrap(c)
-    return similar(c0, D)
+    return DI.recursive_similar(c0, D)
 end
 
 function translate(::Type{D}, contexts::NTuple{C,DI.Context}) where {D<:Dual,C}
@@ -106,7 +106,7 @@ function _translate_toprep(
 end
 function _translate_toprep(::Type{D}, c::DI.Cache) where {D<:Dual}
     c0 = DI.unwrap(c)
-    return similar(c0, D)
+    return DI.recursive_similar(c0, D)
 end
 
 function translate_toprep(::Type{D}, contexts::NTuple{C,DI.Context}) where {D<:Dual,C}

DifferentiationInterface/src/utils/linalg.jl

@@ -10,3 +10,15 @@ At the moment, this only returns `false` for `StaticArrays.SArray`.
 """
 ismutable_array(::Type) = true
 ismutable_array(x) = ismutable_array(typeof(x))
+
+"""
+    recursive_similar(x, T)
+
+Apply `similar(_, T)` recursively to `x` or its components.
+
+Works if `x` is an `AbstractArray` or a (nested) `NTuple` / `NamedTuple` of `AbstractArray`s.
+"""
+recursive_similar(x::AbstractArray, ::Type{T}) where {T} = similar(x, T)
+function recursive_similar(x::Union{Tuple,NamedTuple}, ::Type{T}) where {T}
+    return map(xi -> recursive_similar(xi, T), x)
+end

DifferentiationInterface/test/Core/Internals/linalg.jl

@@ -0,0 +1,9 @@
+using DifferentiationInterface: recursive_similar
+using Test
+
+@test recursive_similar(ones(Int, 2), Float32) isa Vector{Float32}
+@test recursive_similar((ones(Int, 2), ones(Bool, 3, 4)), Float32) isa
+    Tuple{Vector{Float32},Matrix{Float32}}
+@test recursive_similar((a=ones(Int, 2), b=(ones(Bool, 3, 4),)), Float32) isa
+    @NamedTuple{a::Vector{Float32}, b::Tuple{Matrix{Float32}}}
+@test_throws MethodError recursive_similar(1, Float32)