[WIP] Use ChainRules

Project.toml

@@ -3,6 +3,7 @@ uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
 version = "0.3.4"
 
 [deps]
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

src/Zygote.jl

@@ -3,6 +3,7 @@ module Zygote
 using LinearAlgebra, Statistics
 using LinearAlgebra: copytri!, AbstractTriangular
 
+using ChainRules
 import ZygoteRules: @adjoint, @adjoint!, AContext, adjoint, _forward
 
 # This flag enables Zygote to grab extra type inference information during

src/compiler/interface.jl

@@ -28,7 +28,41 @@ end
 
 # Wrappers
 
-_forward(f, args...) = _forward(Context(), f, args...)
+# Note: the ZygoteRules.@adjoint directly overloads this, directly.
+function _forward(f, args...)
+  # First we see if ChainRules has an answer
+  res = rrule(f, args...)
+  if res!==nothing
+    y, crule = res
+    return y, conventionalize(crule)
+  else
+    # No answer from ChainRules, must do nontrival AD
+    return _forward(Context(), f, args...)
+  end
+end
+
+# Convert chainrules into zygote rules
+#TODO: The `nothing` first-arg is not correct in the case of closures
+# ChainRules needs to handle this still
+# See https://github.com/JuliaDiff/ChainRulesCore.jl/issues/4
+conventionalize(crule) = ∇ -> (nothing, conventionalize1(crule(∇)))
+function conventionalize(crules::Tuple)
+  function(∇)
+    ntuple(length(crules) + 1) do ii
+      # first arg is nothing. See TODO above.
+      ii==1 && return nothing
+      crule = crules[ii - 1]
+      conventionalize1(crule(∇))
+    end
+  end
+end
+
+conventionalize1(x) = extern(x)
+conventionalize1(::ChainRules.Zero) = nothing  # Zygote is weird in this way.
+# If it wasn't for Zygote instantiating all partials at once, DNE could just be left to error on extern
+conventionalize1(::ChainRules.DNE) = nothing
+
+
 
 tailmemaybe(::Nothing) = nothing
 tailmemaybe(x::Tuple) = Base.tail(x)

test/features.jl

@@ -73,6 +73,10 @@ end
 
 @test gradient(x -> x.re*x.im, 2+3im) == ((re = 3, im = 2),)
 
+# This tests hitting ChainRules, as without ChainRules rrule definition it would be 1
+# TODO: more clear test, that just defines a ChainRule that has a side-effect we can test for
+@test gradient(abs, 0) == (0,)
+
 struct Foo{T}
   a::T
   b::T