Fixes for #67 (#68)

* Fix identity

* Bump patch version

* Tightens testing + fixes bugs

* Cleans up test

Co-Authored-By: Lyndon White <[email protected]>

* Documents behaviour.

Co-Authored-By: Lyndon White <[email protected]>

* Checks to_vec for loud failure

* Uses oftype for conversion

* Documents to_vec better

Co-authored-by: Lyndon White <[email protected]>

Author: willtebbutt

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"]

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

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/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), ȳ, 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

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")

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