make ChainRulesCore dependency into an extension

Author: KristofferC

Project.toml

@@ -7,14 +7,21 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+ChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 ChainRulesCore = "1"
 julia = "1"
 
 [extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [targets]
-test = ["ChainRulesTestUtils", "Documenter", "ForwardDiff"]
+test = ["ChainRulesCore", "ChainRulesTestUtils", "Documenter", "ForwardDiff"]

ext/ChainRulesCoreExt.jl

@@ -0,0 +1,20 @@
+module ChainRulesCoreExt
+
+using ChainRulesCore
+
+import ChangesOfVariables: _with_ladj_on_mapped
+
+# Need to use a type for this, type inference fails when using a pullback
+# closure over YLT in the rrule, resulting in bad performance:
+struct WithLadjOnMappedPullback{YLT} <: Function end
+function (::WithLadjOnMappedPullback{YLT})(thunked_ΔΩ) where YLT
+    ys, ladj = unthunk(thunked_ΔΩ)
+    return NoTangent(), NoTangent(), map(y -> Tangent{YLT}(y, ladj), ys)
+end
+
+function ChainRulesCore.rrule(::typeof(_with_ladj_on_mapped), map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
+    YLT = eltype(y_with_ladj)
+    return _with_ladj_on_mapped(map_or_bc, y_with_ladj), WithLadjOnMappedPullback{YLT}()
+end
+
+end

src/ChangesOfVariables.jl

@@ -9,11 +9,13 @@ transformations).
 """
 module ChangesOfVariables
 
-using ChainRulesCore
 using LinearAlgebra
 using Test
 
 include("with_ladj.jl")
 include("test.jl")
+if !isdefined(Base, :get_extension)
+    include("../ext/ChainRulesCoreExt.jl")
+end
 
 end # module

src/with_ladj.jl

@@ -114,19 +114,6 @@ function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(
 end
 
 
-# Need to use a type for this, type inference fails when using a pullback
-# closure over YLT in the rrule, resulting in bad performance:
-struct WithLadjOnMappedPullback{YLT} <: Function end
-function (::WithLadjOnMappedPullback{YLT})(thunked_ΔΩ) where YLT
-    ys, ladj = unthunk(thunked_ΔΩ)
-    return NoTangent(), NoTangent(), map(y -> Tangent{YLT}(y, ladj), ys)
-end
-
-function ChainRulesCore.rrule(::typeof(_with_ladj_on_mapped), map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
-    YLT = eltype(y_with_ladj)
-    return _with_ladj_on_mapped(map_or_bc, y_with_ladj), WithLadjOnMappedPullback{YLT}()
-end
-
 function with_logabsdet_jacobian(mapped_f::Base.Fix1{<:Union{typeof(map),typeof(broadcast)}}, X)
     map_or_bc = mapped_f.f
     f = mapped_f.x