Add Tracker support

Author: darsnack

Project.toml

@@ -24,6 +24,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]

src/optimise/Optimise.jl

@@ -4,6 +4,7 @@ using Flux
 using MacroTools: @forward
 import Zygote
 import Zygote: Params, gradient
+import Tracker
 using AbstractDifferentiation
 import Optimisers
 import Optimisers: update, update!

src/optimise/gradients.jl

@@ -21,3 +21,6 @@ AD.@primitive pullback_function(ad::ZygoteExplicitBackend, f, xs...) =
 # this is a hack to get around
 # https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/63#issuecomment-1225959150
 AD.gradient(::ZygoteExplicitBackend, f, xs...) = Zygote.gradient(f, xs...)
+
+# this is to work around AD.TrackerBackend only supporting vectors of params
+AD.gradient(::AD.TrackerBackend, f, xs...) = Tracker.withgradient(f, xs...).grad