Remove custom rrule for _with_ladj_on_mapped

Tricky to implement a correct rrule here that handles
tangents which contain NoTangent or ZeroTangent.

Author: oschulz

Project.toml

@@ -3,25 +3,15 @@ uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 version = "0.1.6"
 
 [deps]
-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]
-ChangesOfVariablesChainRulesCoreExt = "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 = ["ChainRulesCore", "ChainRulesTestUtils", "Documenter", "ForwardDiff"]
+test = ["Documenter", "ForwardDiff"]

README.md

@@ -12,7 +12,7 @@ changes for functions that perform a change of variables (like coordinate
 transformations).
 
 `ChangesOfVariables` is a very lightweight package and has no dependencies
-beyond `Base`, `LinearAlgebra`, `Test` and `ChainRulesCore`.
+beyond `Base`, `LinearAlgebra`, `Test`.
 
 ## Documentation
 

ext/ChangesOfVariablesChainRulesCoreExt.jl

@@ -14,8 +14,5 @@ using Test
 
 include("with_ladj.jl")
 include("test.jl")
-if !isdefined(Base, :get_extension)
-    include("../ext/ChangesOfVariablesChainRulesCoreExt.jl")
-end
 
 end # module

src/ChangesOfVariables.jl

@@ -107,9 +107,13 @@ function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj::Tuple{Any,Real})  where
     return y_with_ladj
 end
 
+_get_all_first(x) = map(first, x)
+# Use x -> x[2] instead of last, using last causes horrible performance in Zygote here:
+_sum_over_second(x) = sum(x -> x[2], x)
+
 function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
-    y = map_or_bc(first, y_with_ladj)
-    ladj = sum(last, y_with_ladj)
+    y = _get_all_first(y_with_ladj)
+    ladj = _sum_over_second(y_with_ladj)
     (y, ladj)
 end
 

src/with_ladj.jl

@@ -7,7 +7,6 @@ using LinearAlgebra
 
 using ChangesOfVariables
 using ChangesOfVariables: test_with_logabsdet_jacobian
-using ChainRulesTestUtils
 
 include("getjacobian.jl")
 
@@ -66,10 +65,4 @@ include("getjacobian.jl")
             test_with_logabsdet_jacobian(f, x, getjacobian)
         end
     end
-
-    @testset "rrules" begin
-        for map_or_bc in (map, broadcast)
-            test_rrule(ChangesOfVariables._with_ladj_on_mapped, map_or_bc, [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)])
-        end
-    end
 end