Backport #770 to ForwardDiff 0.10 (#773)

* Bump actions/cache from 1 to 3 (#643)

Bumps [actions/cache](https://github.com/actions/cache) from 1 to 3.
- [Release notes](https://github.com/actions/cache/releases)
- [Changelog](https://github.com/actions/cache/blob/main/RELEASES.md)
- [Commits](actions/cache@v1...v3)

---
updated-dependencies:
- dependency-name: actions/cache
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Fix tests by using tag types instead of instances (#768)

* Fix gradient and Jacobian for functions with `Dual` output (#770)

* Fix gradient and Jacobian for functions with `Dual` output

* Bump version from 1.2.0 to 1.2.1

* Make tests compatible with Julia 1.0 and 1.6

* Bump version to 0.10.39

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

Author: devmotion

.github/workflows/ci.yml

@@ -32,7 +32,7 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v1
+      - uses: actions/cache@v3
         env:
           cache-name: cache-artifacts
         with:

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "0.10.38"
+version = "0.10.39"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"

ext/ForwardDiffStaticArraysExt.jl

@@ -47,7 +47,7 @@ ForwardDiff._lyap_div!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDi
     result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...)
     return quote
         $(Expr(:meta, :inline))
-        V = StaticArrays.similar_type(S, valtype($y))
+        V = StaticArrays.similar_type(S, valtype(T, $y))
         return V($result)
     end
 end
@@ -76,13 +76,13 @@ end
     result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
     return quote
         $(Expr(:meta, :inline))
-        V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
+        V = StaticArrays.similar_type(S, valtype(T, eltype($ydual)), Size($M, $N))
         return V($result)
     end
 end
 
 function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
-    result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
+    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
     return extract_jacobian!(T, result, ydual, length(x))
 end
 

src/dual.jl

@@ -128,6 +128,17 @@ end
 @inline valtype(::Dual{T,V,N}) where {T,V,N} = V
 @inline valtype(::Type{Dual{T,V,N}}) where {T,V,N} = V
 
+@inline valtype(::Type{T}, ::V) where {T,V} = valtype(T, V)
+@inline valtype(::Type, ::Type{V}) where {V} = V
+@inline valtype(::Type{T}, ::Type{Dual{T,V,N}}) where {T,V,N} = V
+@inline function valtype(::Type{T}, ::Type{Dual{S,V,N}}) where {T,S,V,N}
+    if S ≺ T
+        Dual{S,V,N}
+    else
+        throw(DualMismatchError(T,S))
+    end
+end
+
 @inline tagtype(::V) where {V} = Nothing
 @inline tagtype(::Type{V}) where {V} = Nothing
 @inline tagtype(::Dual{T,V,N}) where {T,V,N} = T

src/gradient.jl

@@ -90,7 +90,7 @@ const GRAD_ERROR = DimensionMismatch("gradient(f, x) expects that f(x) is a real
 function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa Real || throw(GRAD_ERROR)
-    result = similar(x, valtype(ydual))
+    result = similar(x, valtype(T, ydual))
     return extract_gradient!(T, result, ydual)
 end
 
@@ -149,7 +149,7 @@ function chunk_mode_gradient_expr(result_definition::Expr)
 end
 
 @eval function chunk_mode_gradient(f::F, x, cfg::GradientConfig{T,V,N}) where {F,T,V,N}
-    $(chunk_mode_gradient_expr(:(result = similar(x, valtype(ydual)))))
+    $(chunk_mode_gradient_expr(:(result = similar(x, valtype(T, ydual)))))
 end
 
 @eval function chunk_mode_gradient!(result, f::F, x, cfg::GradientConfig{T,V,N}) where {F,T,V,N}

src/jacobian.jl

@@ -128,7 +128,7 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T}
     N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa AbstractArray || throw(JACOBIAN_ERROR)
-    result = similar(ydual, valtype(eltype(ydual)), length(ydual), N)
+    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), N)
     extract_jacobian!(T, result, ydual, N)
     extract_value!(T, result, ydual)
     return result
@@ -217,7 +217,7 @@ end
                                    seed!(xdual, x)
                                end,
                                :(ydual = f(xdual)),
-                               :(result = similar(ydual, valtype(eltype(ydual)), length(ydual), xlen)),
+                               :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)),
                                :()))
 end
 

test/DualTest.jl

@@ -11,6 +11,11 @@ using DiffTests
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
+struct TestTag end
+struct OuterTestTag end
+ForwardDiff.:≺(::Type{TestTag}, ::Type{OuterTestTag}) = true
+ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{<:Tag}) = true
+
 ##################
 # hardcoded test #
 ##################
@@ -179,4 +184,30 @@ end
     @test_throws DomainError ForwardDiff.gradient(x -> x[1]^x[2], [-1.0, 0.5])
 end
 
+# issue #769
+@testset "functions with `Dual` output" begin
+    x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
+    f(x) = sum(ForwardDiff.value, x)
+    der = ForwardDiff.derivative(ForwardDiff.value, first(x))
+
+    # Vector mode
+    grad = ForwardDiff.gradient(f, x)
+    @test grad isa Vector{typeof(der)}
+    @test grad == [der]
+    grad = ForwardDiff.gradient(f, SVector{1}(x))
+    @test grad isa SVector{1,typeof(der)}
+    @test grad == SVector{1}(der)
+
+    # Chunk mode
+    y = repeat(x, 3)
+    cfg = ForwardDiff.GradientConfig(f, y, ForwardDiff.Chunk{2}())
+    grad = ForwardDiff.gradient(f, y, cfg)
+    @test grad isa Vector{typeof(der)}
+    @test grad == [der, der, der]
+    cfg = ForwardDiff.GradientConfig(f, SVector{3}(y), ForwardDiff.Chunk{2}())
+    grad = ForwardDiff.gradient(f, SVector{3}(y), cfg)
+    @test grad isa SVector{3,typeof(der)}
+    @test grad == SVector{3}(der, der, der)
+end
+
 end # module

test/GradientTest.jl

@@ -11,6 +11,11 @@ using LinearAlgebra
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
+struct TestTag end
+struct OuterTestTag end
+ForwardDiff.:≺(::Type{TestTag}, ::Type{OuterTestTag}) = true
+ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{<:Tag}) = true
+
 ##################
 # hardcoded test #
 ##################
@@ -255,4 +260,30 @@ end
     @inferred ForwardDiff.jacobian(g!, [1.0], [0.0])
 end
 
+# issue #769
+@testset "functions with `Dual` output" begin
+    x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
+    f(x) = map(ForwardDiff.value, x)
+    der = ForwardDiff.derivative(ForwardDiff.value, first(x))
+
+    # Vector mode
+    jac = ForwardDiff.jacobian(f, x)
+    @test jac isa Matrix{typeof(der)}
+    @test jac == fill(der, 1, 1)
+    jac = ForwardDiff.jacobian(f, SVector{1}(x))
+    @test jac isa SMatrix{1,1,typeof(der)}
+    @test jac == SMatrix{1,1}(der)
+
+    # Chunk mode
+    y = repeat(x, 3)
+    cfg = ForwardDiff.JacobianConfig(f, y, ForwardDiff.Chunk{2}())
+    jac = ForwardDiff.jacobian(f, y, cfg)
+    @test jac isa Matrix{typeof(der)}
+    @test jac == Diagonal([der, der, der])
+    cfg = ForwardDiff.JacobianConfig(f, SVector{3}(y), ForwardDiff.Chunk{2}())
+    jac = ForwardDiff.jacobian(f, SVector{3}(y), cfg)
+    @test jac isa SMatrix{3,3,typeof(der)}
+    @test jac == Diagonal([der, der, der])
+end
+
 end # module

test/JacobianTest.jl

@@ -152,7 +152,7 @@ intrand(V) = V == Int ? rand(2:10) : rand(V)
 for N in (0,3), V in (Int, Float32), I in (Irrational, AbstractIrrational)
     PARTIALS = ForwardDiff.Partials{N,V}(ntuple(n -> intrand(V), N))
     PRIMAL = intrand(V)
-    FDNUM = ForwardDiff.Dual{TestTag()}(PRIMAL, PARTIALS)
+    FDNUM = ForwardDiff.Dual{TestTag}(PRIMAL, PARTIALS)
 
     @test promote_rule(typeof(FDNUM), I) == promote_rule(I, typeof(FDNUM))
     # π::Irrational, twoπ::AbstractIrrational