Add documentation

Author: oschulz

README.md

@@ -6,6 +6,11 @@
 [![Build Status](https://github.com/oschulz/ForwardDiffPullbacks.jl/workflows/CI/badge.svg?branch=master)](https://github.com/oschulz/ForwardDiffPullbacks.jl/actions?query=workflow%3ACI)
 [![Codecov](https://codecov.io/gh/oschulz/ForwardDiffPullbacks.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/oschulz/ForwardDiffPullbacks.jl)
 
+ForwardDiffPullbacks implements pullbacks compatible with
+[ChainRulesCore](https://github.com/JuliaDiff/ChainRulesCore.jl) that are calculated via
+[ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl).
+
+See the documentation for details.
 
 ## Documentation
 

docs/src/index.md

@@ -1 +1,9 @@
 # ForwardDiffPullbacks.jl
+
+ForwardDiffPullbacks implements pullbacks compatible with [ChainRulesCore](https://github.com/JuliaDiff/ChainRulesCore.jl) that are calculated via [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl).
+
+This package provides the function [`fwddiff`](@ref). If wrapped around a function (i.e. `fwddiff(f)`), it will cause ChainRules (and implicitly Zygote) pullbacks to be calculated using ForwardDiff (i.e. by evaluating the original function with `ForwardDiff.Dual` numbers, possibly multiple times). The pullback will return a ChainRule thunk for each argument of the function.
+
+So `Zygote.gradient(fwddiff(f), xs...)` should yield the same result as `Zygote.gradient(f, xs...)`, but will typically be substantially faster a function that has a comparatively small number of arguments, especially if the function runs a deep calculation. Broadcasting (i.e. `g.(fwddiff(f))`) is supported as well.
+
+Currently, ForwardDiffPullbacks supports functions with `Real`, `Tuple` and `StaticArrays.SVector` arguments. Support for `StaticArrays.SArray` and `Array`-valued arguments in general is on the to-do list.

src/with_forwarddiff.jl

@@ -10,13 +10,41 @@ end
 # Desireable for consistent behavior?
 # Base.broadcasted(wrapped_f::WithForwardDiff, xs...) = broadcast(wrapped_f.f, xs...)
 
+
 """
     fwddiff(f::Base.Callable)::Function
 
-Use `ForwardDiff` in `ChainRulesCore` pullback For
+Use `ForwardDiff` dual numbers to implement `ChainRulesCore` pullbacks For
 
 * `fwddiff(f)(args...)
 * `fwddiff(f).(args...)
+
+Example:
+
+```
+using ForwardDiffPullbacks, StaticArrays
+
+f = (xs...) -> (sum(map(x -> sum(map(x -> x^2, x)), xs)))
+xs = (2, (3, 4), SVector(5, 6, 7))
+f(xs...) == 139
+
+using ChainRulesCore
+
+y, back = rrule(fwddiff(f), xs...)
+y == 139
+map(unthunk, back(1)) == (Zero(), 4, (6, 8), [10, 12, 14])
+
+using Zygote
+
+Zygote.gradient(fwddiff(f), xs...) == Zygote.gradient(f, xs...)
+
+Xs = map(x -> fill(x, 100), xs)
+Zygote.gradient((Xs...) -> sum(fwddiff(f).(Xs...)), Xs...) ==
+    Zygote.gradient((Xs...) -> sum(f.(Xs...)), Xs...)
+```
+
+The gradient is the same with and without `fwddiff`, but `fwddiff` makes the
+gradient calculation a lot faster here.
 """
 fwddiff(f::Base.Callable) = WithForwardDiff(f)
 export fwddiff