Repository: JuliaDiff/FiniteDifferences.jl
Pull request #68 (2020-04-08 23:08:04)

Fixes for #67
This is a partial resolution for #67 @sethaxen - there's the still the `StaticArrays` issue, that requires a bit more work.

I'm not sure if this is going to the be correct solution forever, but it's a decent patch for now.

willtebbutt: This now fixes all of #67 , and I'm pretty confident that the changes aren't breaking. There are a bunch of extra tests to make sure that the `from_vec` bit of `to_vec` returns precisely the right type, and to ensure that `SArray`s play nicely with `to_vec`.

These extra tests threw up a couple of bugs that (somehow) hadn't been causing problems, but were definitely bugs. These bugs have been fixed.
sethaxen: > @sethaxen I would appreciate it if you could test this branch against Manifolds to check that it's fixed your problems? If not I'll figure out what other tests we should be adding.

I tested this branch against our master, and the test no longer fails. Thanks for the quick fix!
willtebbutt: Thanks for checking!

@oxinabox I've implemented your `oftype` suggestions, but didn't merge your suggestions directly because there were related changes that needed to happen for code-cleanliness as a result.

Patch:
diff --git a/Project.toml b/Project.toml
index 7388759..b4a2a38 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "FiniteDifferences"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.9.3"
+version = "0.9.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -12,6 +12,7 @@ julia = "1"
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [targets]
-test = ["Random", "Test"]
+test = ["Random", "StaticArrays", "Test"]
diff --git a/src/grad.jl b/src/grad.jl
index 80f6998..62ac4cd 100644
--- a/src/grad.jl
+++ b/src/grad.jl
@@ -16,7 +16,7 @@ function jacobian(fdm, f, x::Vector{<:Number}; len=nothing)
         return fdm(zero(eltype(x))) do ε
             xn = x[n]
             x[n] = xn + ε
-            ret = first(to_vec(f(x)))
+            ret = copy(first(to_vec(f(x))))  # copy required incase `f(x)` returns something that aliases `x`
             x[n] = xn  # Can't do `x[n] -= ϵ` as floating-point math is not associative
             return ret
         end
diff --git a/src/to_vec.jl b/src/to_vec.jl
index be54748..d0c23aa 100644
--- a/src/to_vec.jl
+++ b/src/to_vec.jl
@@ -6,40 +6,38 @@ transformation.
 """
 function to_vec(x::Number)
     function Number_from_vec(x_vec)
-        return first(x_vec)
+        return oftype(x, first(x_vec))
     end
     return [x], Number_from_vec
 end
 
-# Vectors
+# Base case -- if x is already a Vector{<:Number} there's no conversion necessary.
 to_vec(x::Vector{<:Number}) = (x, identity)
-function to_vec(x::Vector)
+
+function to_vec(x::AbstractVector)
     x_vecs_and_backs = map(to_vec, x)
     x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
     function Vector_from_vec(x_vec)
         sz = cumsum(map(length, x_vecs))
-        return [backs[n](x_vec[sz[n] - length(x_vecs[n]) + 1:sz[n]]) for n in eachindex(x)]
+        x_Vec = [backs[n](x_vec[sz[n] - length(x_vecs[n]) + 1:sz[n]]) for n in eachindex(x)]
+        return oftype(x, x_Vec)
     end
     return vcat(x_vecs...), Vector_from_vec
 end
 
-# Arrays
-function to_vec(x::Array{<:Number})
-    function Array_from_vec(x_vec)
-        return reshape(x_vec, size(x))
-    end
-    return vec(x), Array_from_vec
-end
+function to_vec(x::AbstractArray)
+
+    x_vec, from_vec = to_vec(vec(x))
 
-function to_vec(x::Array)
-    x_vec, back = to_vec(reshape(x, :))
     function Array_from_vec(x_vec)
-        return reshape(back(x_vec), size(x))
+        return oftype(x, reshape(from_vec(x_vec), size(x)))
     end
+
     return x_vec, Array_from_vec
 end
 
-# AbstractArrays
+# Some specific subtypes of AbstractArray.
+
 function to_vec(x::T) where {T<:LinearAlgebra.AbstractTriangular}
     x_vec, back = to_vec(Matrix(x))
     function AbstractTriangular_from_vec(x_vec)
@@ -63,17 +61,25 @@ function to_vec(X::Diagonal)
 end
 
 function to_vec(X::Transpose)
+
+    x_vec, x_from_vec = to_vec(X.parent)
+
     function Transpose_from_vec(x_vec)
-        return Transpose(permutedims(reshape(x_vec, size(X))))
+        return Transpose(x_from_vec(x_vec))
     end
-    return vec(Matrix(X)), Transpose_from_vec
+
+    return x_vec, Transpose_from_vec
 end
 
 function to_vec(X::Adjoint)
+
+    x_vec, x_from_vec = to_vec(X.parent)
+
     function Adjoint_from_vec(x_vec)
-        return Adjoint(conj!(permutedims(reshape(x_vec, size(X)))))
+        return Adjoint(x_from_vec(x_vec))
     end
-    return vec(Matrix(X)), Adjoint_from_vec
+
+    return x_vec, Adjoint_from_vec
 end
 
 # Non-array data structures


Test patch:
diff --git a/test/grad.jl b/test/grad.jl
index fa0607f..e47667a 100644
--- a/test/grad.jl
+++ b/test/grad.jl
@@ -52,6 +52,9 @@ using FiniteDifferences: grad, jacobian, _jvp, jvp, j′vp, _j′vp, to_vec
         @test Ac == A
         check_jac_and_jvp_and_j′vp(fdm, x->sin.(A * x), ȳ, x, ẋ, cos.(A * x) .* A)
         @test Ac == A
+
+        # Prevent regression against https://github.com/JuliaDiff/FiniteDifferences.jl/issues/67
+        @test first(jacobian(fdm, identity, x)) ≈ one(Matrix{T}(undef, length(x), length(x)))
     end
 
     @testset "multi vars jacobian/grad" begin
diff --git a/test/runtests.jl b/test/runtests.jl
index 70f9fd3..485e194 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,4 +1,4 @@
-using FiniteDifferences, Test, Random, Printf, LinearAlgebra
+using FiniteDifferences, Test, Random, Printf, LinearAlgebra, StaticArrays
 
 @testset "FiniteDifferences" begin
     include("methods.jl")
diff --git a/test/to_vec.jl b/test/to_vec.jl
index 967ec8a..229746f 100644
--- a/test/to_vec.jl
+++ b/test/to_vec.jl
@@ -11,15 +11,26 @@ end
 Base.:(==)(x::DummyType, y::DummyType) = x.X == y.X
 Base.length(x::DummyType) = size(x.X, 1)
 
-function test_to_vec(x)
+# A dummy FillVector. This is a type for which the fallback implementation of
+# `to_vec` should fail loudly.
+struct FillVector <: AbstractVector{Float64}
+    x::Float64
+    len::Int
+end
+
+Base.size(x::FillVector) = (x.len,)
+Base.getindex(x::FillVector, n::Int) = x.x
+
+function test_to_vec(x::T) where {T}
     x_vec, back = to_vec(x)
     @test x_vec isa Vector
     @test x == back(x_vec)
+    @test back(x_vec) isa T
     return nothing
 end
 
 @testset "to_vec" begin
-    @testset "$T" for T in (Float64, ComplexF64)
+    @testset "$T" for T in (Float32, ComplexF32, Float64, ComplexF64)
         if T == Float64
             test_to_vec(1.0)
             test_to_vec(1)
@@ -38,6 +49,8 @@ end
         test_to_vec(Symmetric(randn(T, 11, 11)))
         test_to_vec(Diagonal(randn(T, 7)))
         test_to_vec(DummyType(randn(T, 2, 9)))
+        test_to_vec(SVector{2, T}(1.0, 2.0))
+        test_to_vec(SMatrix{2, 2, T}(1.0, 2.0, 3.0, 4.0))
 
         @testset "$Op" for Op in (Adjoint, Transpose)
             test_to_vec(Op(randn(T, 4, 4)))
@@ -62,4 +75,10 @@ end
             end
         end
     end
+
+    @testset "FillVector" begin
+        x = FillVector(5.0, 10)
+        x_vec, from_vec = to_vec(x)
+        @test_throws MethodError from_vec(randn(10))
+    end
 end