Merge pull request #192 from JuliaDiff/ox/projrand

use ProjectTo for rand_tangent, and also catch some more places we shoud NoTangent

Author: oxinabox

Project.toml

@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.10.12, 1"
+ChainRulesCore = "1"
 Compat = "3"
 FiniteDifferences = "0.12.12"
 julia = "1"

src/rand_tangent.jl

@@ -29,19 +29,21 @@ end
 # multiply by 9 to give a bigger range of values tested: no so tightly clustered around 0.
 rand_tangent(rng::AbstractRNG, ::BigFloat) = round(big(9 * randn(rng)), sigdigits=5, base=2)
 
-rand_tangent(rng::AbstractRNG, x::StridedArray{T, 0}) where {T} = fill(rand_tangent(x[1]))
-rand_tangent(rng::AbstractRNG, x::StridedArray) = rand_tangent.(Ref(rng), x)
-rand_tangent(rng::AbstractRNG, x::Adjoint) = adjoint(rand_tangent(rng, parent(x)))
-rand_tangent(rng::AbstractRNG, x::Transpose) = transpose(rand_tangent(rng, parent(x)))
 
-function rand_tangent(rng::AbstractRNG, x::T) where {T<:Tuple}
-    return Tangent{T}(rand_tangent.(Ref(rng), x)...)
-end
+rand_tangent(rng::AbstractRNG, x::Array{<:Any, 0}) = _compress_notangent(fill(rand_tangent(rng, x[])))
+rand_tangent(rng::AbstractRNG, x::Array) = _compress_notangent(rand_tangent.(Ref(rng), x))
 
-function rand_tangent(rng::AbstractRNG, xs::T) where {T<:NamedTuple}
-    return Tangent{T}(; map(x -> rand_tangent(rng, x), xs)...)
+# All other AbstractArray's can be handled using the ProjectTo mechanics.
+# and follow the same requirements
+function rand_tangent(rng::AbstractRNG, x::AbstractArray)
+    return _compress_notangent(ProjectTo(x)(rand_tangent(rng, collect(x))))
 end
 
+# TODO: arguably ProjectTo should handle this for us for AbstactArrays
+# https://github.com/JuliaDiff/ChainRulesCore.jl/issues/410
+_compress_notangent(::AbstractArray{NoTangent}) = NoTangent()
+_compress_notangent(x) = x
+
 function rand_tangent(rng::AbstractRNG, x::T) where {T}
     if !isstructtype(T)
         throw(ArgumentError("Non-struct types are not supported by this fallback."))
@@ -54,8 +56,12 @@ function rand_tangent(rng::AbstractRNG, x::T) where {T}
     if all(tangent isa NoTangent for tangent in tangents)
         # if none of my fields can be perturbed then I can't be perturbed
         return NoTangent()
+    end
+
+    if T <: Tuple
+        return Tangent{T}(tangents...)
     else
-        Tangent{T}(; NamedTuple{field_names}(tangents)...)
+        return Tangent{T}(; NamedTuple{field_names}(tangents)...)
     end
 end
 

test/rand_tangent.jl

@@ -39,44 +39,54 @@ struct Bar
         (randn(Complex{Float32}, 5, 4), Matrix{Complex{Float32}}),
         ([randn(5, 4), 4.0], Vector{Any}),
 
-        # Wrapper Arrays
-        (randn(5, 4)', Adjoint{Float64, Matrix{Float64}}),
-        (transpose(randn(5, 4)), Transpose{Float64, Matrix{Float64}}),
-
+        # Co-Arrays
+        (randn(5)', Adjoint{Float64, Vector{Float64}}),  # row-vector: special
+        (randn(5, 4)', Matrix{Float64}),                 # matrix: generic dense
+        
+        (transpose(randn(5)), Transpose{Float64, Vector{Float64}}),  # row-vector: special
+        (transpose(randn(5, 4)), Matrix{Float64}),                   # matrix: generic dense
+        
+        # AbstactArrays of non-perturbable types
+        (1:10, NoTangent),
+        (1:2:10, NoTangent),
+        ([false, true], NoTangent),
 
         # Tuples.
         ((4.0, ), Tangent{Tuple{Float64}}),
         ((5.0, randn(3)), Tangent{Tuple{Float64, Vector{Float64}}}),
+        ((false, true), NoTangent),
+        (Tuple{}(), NoTangent),
 
         # NamedTuples.
         ((a=4.0, ), Tangent{NamedTuple{(:a,), Tuple{Float64}}}),
         ((a=5.0, b=1), Tangent{NamedTuple{(:a, :b), Tuple{Float64, Int}}}),
+        ((a=false, b=true), NoTangent),
+        ((;), NoTangent),
 
         # structs.
         (Bar(5.0, 4, rand(rng, 3)), Tangent{Bar}),
         (Bar(4.0, 3, Bar(5.0, 2, 4)), Tangent{Bar}),
         (sin, NoTangent),
         # all fields NoTangent implies NoTangent
         (Pair(:a, "b"), NoTangent),
-        (1:10, NoTangent),
-        (1:2:10, NoTangent),
+        (CartesianIndex(2, 3), NoTangent),
 
-        # LinearAlgebra types (also just structs).
+        # LinearAlgebra types
         (
             UpperTriangular(randn(3, 3)),
-            Tangent{UpperTriangular{Float64, Matrix{Float64}}},
+            UpperTriangular{Float64, Matrix{Float64}},
         ),
         (
             Diagonal(randn(2)),
-            Tangent{Diagonal{Float64, Vector{Float64}}},
+            Diagonal{Float64, Vector{Float64}},
         ),
         (
             Symmetric(randn(2, 2)),
-            Tangent{Symmetric{Float64, Matrix{Float64}}},
+            Symmetric{Float64, Matrix{Float64}},
         ),
         (
             Hermitian(randn(ComplexF64, 1, 1)),
-            Tangent{Hermitian{ComplexF64, Matrix{ComplexF64}}},
+            Hermitian{ComplexF64, Matrix{ComplexF64}},
         ),
     ]
         @test rand_tangent(rng, x) isa T_tangent
@@ -96,6 +106,7 @@ struct Bar
 
     # Julia 1.6 changed to using Ryu printing algorithm and seems better at printing short
     VERSION > v"1.6" && @testset "niceness of printing" begin
+        rng = MersenneTwister()
         for i in 1:50
             @test length(string(rand_tangent(1.0))) <= 6
             @test length(string(rand_tangent(1.0 + 1.0im))) <= 12
@@ -104,3 +115,4 @@ struct Bar
         end
     end
 end
+

test/runtests.jl

@@ -1,5 +1,6 @@
 using ChainRulesCore
 using ChainRulesTestUtils
+using ChainRulesTestUtils: rand_tangent
 using FiniteDifferences
 using LinearAlgebra
 using Random