Automatic tangent generation for simple types

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.9"
+ChainRulesCore = "0.9.1"
 Compat = "3"
 FiniteDifferences = "0.10"
 julia = "1"

src/ChainRulesTestUtils.jl

@@ -13,6 +13,7 @@ const _fdm = central_fdm(5, 1)
 
 export test_scalar, frule_test, rrule_test, generate_well_conditioned_matrix
 
+include("generate_tangent.jl")
 include("to_vec.jl")
 include("isapprox.jl")
 include("data_generation.jl")

src/generate_tangent.jl

@@ -0,0 +1,40 @@
+"""
+    rand_tangent([rng::AbstractRNG,] x)
+
+Returns a randomly generated tangent vector appropriate for the primal value `x`.
+"""
+rand_tangent(x) = rand_tangent(Random.default_rng(), x)
+
+function rand_tangent(rng::AbstractRNG, x::Union{Symbol, AbstractChar, AbstractString})
+    return DoesNotExist()
+end
+
+rand_tangent(rng::AbstractRNG, x::Integer) = DoesNotExist()
+
+rand_tangent(rng::AbstractRNG, x::T) where {T<:Number} = randn(rng, T)
+
+rand_tangent(rng::AbstractRNG, x::StridedArray) = rand_tangent.(Ref(rng), x)
+
+function rand_tangent(rng::AbstractRNG, x::T) where {T<:Tuple}
+    return Composite{T}(rand_tangent.(Ref(rng), x)...)
+end
+
+function rand_tangent(rng::AbstractRNG, xs::T) where {T<:NamedTuple}
+    return Composite{T}(; map(x -> rand_tangent(rng, x), xs)...)
+end
+
+function rand_tangent(rng::AbstractRNG, x::T) where {T}
+    if !isstructtype(T)
+        throw(ArgumentError("Non-struct types are not supported by this fallback."))
+    end
+
+    field_names = fieldnames(T)
+    if length(field_names) > 0
+        tangents = map(field_names) do field_name
+            rand_tangent(rng, getfield(x, field_name))
+        end
+        return Composite{T}(; NamedTuple{field_names}(tangents)...)
+    else
+        return NO_FIELDS
+    end
+end

test/generate_tangent.jl

@@ -0,0 +1,41 @@
+using ChainRulesTestUtils: rand_tangent
+
+# Test struct for `rand_tangent`.
+struct Foo
+   a::Float64
+   b::Int
+   c::Any
+end
+
+@testset "generate_tangent" begin
+    rng = MersenneTwister(123456)
+
+    foreach([
+        ("hi", DoesNotExist),
+        ('a', DoesNotExist),
+        (:a, DoesNotExist),
+        (true, DoesNotExist),
+        (4, DoesNotExist),
+        (5.0, Float64),
+        (5.0 + 0.4im, Complex{Float64}),
+        (randn(Float32, 3), Vector{Float32}),
+        (randn(Complex{Float64}, 2), Vector{Complex{Float64}}),
+        (randn(5, 4), Matrix{Float64}),
+        (randn(Complex{Float32}, 5, 4), Matrix{Complex{Float32}}),
+        ([randn(5, 4), 4.0], Vector{Any}),
+        ((4.0, ), Composite{Tuple{Float64}}),
+        ((5.0, randn(3)), Composite{Tuple{Float64, Vector{Float64}}}),
+        ((a=4.0, ), Composite{NamedTuple{(:a,), Tuple{Float64}}}),
+        ((a=5.0, b=1), Composite{NamedTuple{(:a, :b), Tuple{Float64, Int}}}),
+        (sin, typeof(NO_FIELDS)),
+        (Foo(5.0, 4, rand(rng, 3)), Composite{Foo}),
+        (Foo(4.0, 3, Foo(5.0, 2, 4)), Composite{Foo}),
+    ]) do (x, T_tangent)
+        @test rand_tangent(rng, x) isa T_tangent
+        @test rand_tangent(x) isa T_tangent
+        @test x + rand_tangent(rng, x) isa typeof(x)
+    end
+
+    # Ensure struct fallback errors for non-struct types.
+    @test_throws ArgumentError invoke(rand_tangent, Tuple{AbstractRNG, Any}, rng, 5.0)
+end

test/runtests.jl

@@ -5,6 +5,7 @@ using Random
 using Test
 
 @testset "ChainRulesTestUtils.jl" begin
+    include("generate_tangent.jl")
     include("to_vec.jl")
     include("isapprox.jl")
     include("testers.jl")