Merge pull request #279 from JuliaDiff/ox/inplacecheck

Correct inplace for complex BLAS and fix some accumulation against nondifferentails

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.7.24"
+version = "0.7.25"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -12,9 +12,9 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 ChainRulesCore = "0.9.12"
-ChainRulesTestUtils = "0.4.2, 0.5"
+ChainRulesTestUtils = "0.5"
 Compat = "3"
-FiniteDifferences = "0.10"
+FiniteDifferences = "0.11"
 Reexport = "0.2"
 Requires = "0.5.2, 1"
 julia = "1"

src/rulesets/LinearAlgebra/blas.jl

@@ -129,7 +129,7 @@ function rrule(::typeof(gemv), tA::Char, α::T, A::AbstractMatrix{T},
         else  # uppercase(tA) === 'T'
             ∂A = InplaceableThunk(
                 @thunk(conj(α * x * ȳ')),
-                Ā -> conj!(ger!(α, x, ȳ, Ā))
+                Ā -> conj!(ger!(α, x, ȳ, conj!(Ā)))
             )
             ∂x = InplaceableThunk(
                 @thunk(gemv('N', α', conj(A), ȳ)),
@@ -188,7 +188,7 @@ function rrule(
                 )
                 ∂B = InplaceableThunk(
                     @thunk(conj(gemm('C', 'N', α, C̄, A))),
-                    B̄ -> conj!(gemm!('C', 'N', α, C̄, A, β, B̄))
+                    B̄ -> conj!(gemm!('C', 'N', α, C̄, A, β, conj!(B̄)))
                 )
             end
         elseif uppercase(tA) === 'C'
@@ -224,7 +224,7 @@ function rrule(
             if uppercase(tB) === 'N'
                 ∂A = InplaceableThunk(
                     @thunk(conj(gemm('N', 'C', α, B, C̄))),
-                    Ā -> conj!(gemm!('N', 'C', α, B, C̄, β, Ā))
+                    Ā -> conj!(gemm!('N', 'C', α, B, C̄, β, conj!(Ā)))
                 )
                 ∂B = InplaceableThunk(
                     @thunk(gemm('N', 'N', α', conj(A), C̄)),

test/rulesets/Base/base.jl

@@ -43,7 +43,7 @@
             test_scalar(acscd, 1/x)
             test_scalar(acotd, 1/x)
         end
-        
+
         @testset "sinc" for x = (0.0, 0.434, Complex(0.434, 0.25))
             test_scalar(sinc, x)
         end
@@ -126,14 +126,14 @@
         frule_test(identity, (randn(T, 4), randn(T, 4)))
         frule_test(
             identity,
-            (Composite{Tuple}(randn(T, 3)...), Composite{Tuple}(randn(T, 3)...))
+            (Tuple(randn(T, 3)), Composite{Tuple{T, T, T}}(randn(T, 3)...))
         )
 
         rrule_test(identity, randn(T), (randn(T), randn(T)))
         rrule_test(identity, randn(T, 4), (randn(T, 4), randn(T, 4)))
         rrule_test(
-            identity, Tuple(randn(T, 3)),
-            (Composite{Tuple}(randn(T, 3)...), Composite{Tuple}(randn(T, 3)...))
+            identity, Composite{Tuple{T, T, T}}(randn(T, 3)...),
+            (Tuple(randn(T, 3)), Composite{Tuple{T, T, T}}(randn(T, 3)...))
         )
     end
 

test/rulesets/Base/evalpoly.jl

@@ -27,9 +27,8 @@
             Ω = evalpoly(x, p)
             Ω̄ = randn(T, size(Ω)...)
             frule_test(evalpoly, (x, ẋ), (p, ṗ))
-            frule_test(evalpoly, (x, ẋ), (Tuple(p), Tuple(ṗ)))
+            frule_test(evalpoly, (x, ẋ), (Tuple(p), rand_tangent(Tuple(p))))
             rrule_test(evalpoly, Ω̄, (x, x̄), (p, p̄))
-            rrule_test(evalpoly, Ω̄, (x, x̄), (Tuple(p), Tuple(p̄)))
+            rrule_test(evalpoly, Ω̄, (x, x̄), (Tuple(p), rand_tangent(Tuple(p))))
         end
     end
-

test/rulesets/Base/fastmath_able.jl

@@ -38,6 +38,8 @@ function complex_jacobian_test(f, z)
 end
 
 const FASTABLE_AST = quote
+    is_fastmath_mode = sin === Base.FastMath.sin_fast
+
     @testset "Trig" begin
         @testset "Basics" for x = (Float64(π)-0.01, Complex(π, π/2))
             test_scalar(sin, x)
@@ -66,7 +68,7 @@ const FASTABLE_AST = quote
     end
 
     @testset "exponents" begin
-        for x in (-0.1, 6.4, 0.5 + 0.25im)
+        for x in (-0.1, 7.9, 0.5 + 0.25im)
             test_scalar(inv, x)
 
             test_scalar(exp, x)
@@ -117,14 +119,14 @@ const FASTABLE_AST = quote
         for x in (-4.1, 6.4, 0.0, 0.0 + 0.0im, 0.5 + 0.25im)
             test_scalar(+, x)
             test_scalar(-, x)
-            test_scalar(atan, x)
+            test_scalar(atan, x; rtol=(is_fastmath_mode ? 1e-7 : 1e-9))
         end
     end
 
     @testset "binary functions" begin
         @testset "$f(x, y)" for f in (atan, rem, max, min)
-            x, Δx, x̄ = 10rand(3)
-            y, Δy, ȳ = rand(3)
+            x, Δx, x̄ = 100rand(3)
+            y, Δy, ȳ = 10rand(3)
             Δz = rand()
 
             frule_test(f, (x, Δx), (y, Δy))
@@ -136,8 +138,9 @@ const FASTABLE_AST = quote
             y, Δy, ȳ = rand(T, 3)
             Δz = randn(typeof(f(x, y)))
 
-            frule_test(f, (x, Δx), (y, Δy))
-            rrule_test(f, Δz, (x, x̄), (y, ȳ))    
+            # some tests struggle in fast_math mode to get accurasy so we relax it some.
+            frule_test(f, (x, Δx), (y, Δy); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
+            rrule_test(f, Δz, (x, x̄), (y, ȳ); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
         end
 
         @testset "$f(x::$T, y::$T) type check" for f in (/, +, -,\, hypot, ^), T in (Float32, Float64)
@@ -151,11 +154,11 @@ const FASTABLE_AST = quote
             @test extern.((∂x, ∂y)) isa Tuple{T, T}
 
             if f != hypot
-                # Issue #233                
+                # Issue #233
                 @test frule((Zero(), Δx, Δy), f, x, 2) isa Tuple{T, T}
                 _, ∂x, ∂y = rrule(f, x, 2)[2](Δz)
                 @test extern.((∂x, ∂y)) isa Tuple{T, T}
-                
+
                 @test frule((Zero(), Δx, Δy), f, 2, y) isa Tuple{T, T}
                 _, ∂x, ∂y = rrule(f, 2, y)[2](Δz)
                 @test extern.((∂x, ∂y)) isa Tuple{T, T}
@@ -164,17 +167,17 @@ const FASTABLE_AST = quote
 
         @testset "^(x::$T, n::$T)" for T in (Float64, ComplexF64)
             # for real x and n, x must be >0
-            x = T <: Real ? 15rand() : 15randn(ComplexF64)
-            Δx, x̄ = 10rand(T, 2)
-            y, Δy, ȳ = rand(T, 3)
+            x, Δx, x̄ = rand(T, 3) .+ 3
+            y, Δy, ȳ = rand(T, 3) .+ 3
             Δz = rand(T)
 
             frule_test(^, (x, Δx), (y, Δy))
             rrule_test(^, Δz, (x, x̄), (y, ȳ))
 
             T <: Real && @testset "discontinuity for ^(x::Real, n::Int) when x ≤ 0" begin
                 # finite differences doesn't work for x < 0, so we check manually
-                x, y = -10rand(T), 3
+                x = -rand(T) .- 3
+                y = 3
                 Δx = randn(T)
                 Δy = randn(T)
                 Δz = randn(T)

test/rulesets/LinearAlgebra/blas.jl

@@ -33,7 +33,7 @@
                 n = 10
                 x, ẋ, x̄ = randn(T, n), randn(T, n), randn(T, n)
                 frule_test(BLAS.nrm2, (x, ẋ))
-                rrule_test(BLAS.nrm2, randn(), (x, x̄))
+                rrule_test(BLAS.nrm2, randn(), (x, x̄); rtol=1e-7)
             end
         end
 

test/rulesets/LinearAlgebra/factorization.jl

@@ -89,10 +89,10 @@ using ChainRules: level2partition, level3partition, chol_blocked_rev, chol_unblo
                 ΔF = unthunk(dF)
                 _, dX = dX_pullback(ΔF)
                 X̄_ad = dot(unthunk(dX), V)
-                X̄_fd = _fdm() do ε
+                X̄_fd = central_fdm(5, 1)(0.000_001) do ε
                     dot(Ȳ, getproperty(cholesky(X .+ ε .* V), p))
                 end
-                @test X̄_ad ≈ X̄_fd rtol=1e-6 atol=1e-6
+                @test X̄_ad ≈ X̄_fd rtol=1e-4
             end
         end
         @testset "helper functions" begin

test/rulesets/packages/SpecialFunctions.jl

@@ -50,7 +50,7 @@ end
             isreal(x) || continue
 
             Δx, x̄ = randn(2)
-            Δz = (randn(), randn())
+            Δz = Composite{Tuple{Float64, DoesNotExist}}(randn(), Zero())
 
             frule_test(SpecialFunctions.logabsgamma, (x, Δx))
             rrule_test(SpecialFunctions.logabsgamma, Δz, (x, x̄))

test/runtests.jl

@@ -5,6 +5,7 @@ using ChainRulesTestUtils
 using ChainRulesTestUtils: _fdm
 using Compat: only
 using FiniteDifferences
+using FiniteDifferences: rand_tangent
 using LinearAlgebra
 using LinearAlgebra.BLAS
 using LinearAlgebra: dot