Move ChainRulesCore support to extension

Author: oschulz

Project.toml

@@ -4,7 +4,6 @@ authors = ["Chad Scherrer <[email protected]>", "Oliver Schulz <oschulz@mp
 version = "0.14.11"
 
 [deps]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -29,6 +28,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+MeasureBaseChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"

ext/MeasureBaseChainRulesCoreExt.jl

@@ -0,0 +1,37 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseChainRulesCoreExt
+
+using MeasureBase
+using ChainRulesCore: NoTangent, ZeroTangent
+import ChainRulesCore
+
+
+@inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
+    y = _checksupport(cond, result)
+    function _checksupport_pullback(ȳ)
+        return NoTangent(), ZeroTangent(), one(ȳ)
+    end
+    y, _checksupport_pullback
+end
+
+
+_require_insupport_pullback(ΔΩ) = NoTangent(), ZeroTangent()
+function ChainRulesCore.rrule(::typeof(require_insupport), μ, x)
+    return require_insupport(μ, x), _require_insupport_pullback
+end
+
+
+_origin_depth_pullback(ΔΩ) = NoTangent(), NoTangent()
+ChainRulesCore.rrule(::typeof(_origin_depth), ν) = _origin_depth(ν), _origin_depth_pullback
+
+
+_check_dof_pullback(ΔΩ) = NoTangent(), NoTangent(), NoTangent()
+ChainRulesCore.rrule(::typeof(check_dof), ν, μ) = check_dof(ν, μ), _check_dof_pullback
+
+
+_checked_arg_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ
+ChainRulesCore.rrule(::typeof(checked_arg), ν, x) = checked_arg(ν, x), _checked_arg_pullback
+
+
+end # module MeasureBaseChainRulesCoreExt

src/MeasureBase.jl

@@ -31,7 +31,6 @@ using IntervalSets
 using PrettyPrinting
 const Pretty = PrettyPrinting
 
-using ChainRulesCore
 import FillArrays
 using Static
 using Static: StaticInteger

src/density-core.jl

@@ -35,14 +35,6 @@ end
 
 _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf))
 
-import ChainRulesCore
-@inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
-    y = _checksupport(cond, result)
-    function _checksupport_pullback(ȳ)
-        return NoTangent(), ZeroTangent(), one(ȳ)
-    end
-    y, _checksupport_pullback
-end
 
 export unsafe_logdensityof
 

src/getdof.jl

@@ -51,6 +51,7 @@ end
 _check_dof_pullback(ΔΩ) = NoTangent(), NoTangent(), NoTangent()
 ChainRulesCore.rrule(::typeof(check_dof), ν, μ) = check_dof(ν, μ), _check_dof_pullback
 
+
 """
     MeasureBase.NoArgCheck{MU,T}
 
@@ -78,6 +79,3 @@ end
 @propagate_inbounds function checked_arg(mu::MU, x) where {MU}
     _default_checked_arg(MU, basemeasure(mu), x)
 end
-
-_checked_arg_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ
-ChainRulesCore.rrule(::typeof(checked_arg), ν, x) = checked_arg(ν, x), _checked_arg_pullback

src/insupport.jl

@@ -18,11 +18,6 @@ Checks if `x` is in the support of distribution/measure `μ`, throws an
 """
 function require_insupport end
 
-_require_insupport_pullback(ΔΩ) = NoTangent(), ZeroTangent()
-function ChainRulesCore.rrule(::typeof(require_insupport), μ, x)
-    return require_insupport(μ, x), _require_insupport_pullback
-end
-
 function require_insupport(μ, x)
     if !insupport(μ, x)
         throw(ArgumentError("x is not within the support of μ"))

src/transport.jl

@@ -135,9 +135,6 @@ end
     return static(10)
 end
 
-_origin_depth_pullback(ΔΩ) = NoTangent(), NoTangent()
-ChainRulesCore.rrule(::typeof(_origin_depth), ν) = _origin_depth(ν), _origin_depth_pullback
-
 # If both both measures have no origin:
 function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x)
     _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x)