docs: add mixed-mode explanations (#690)

Author: gdalle

DifferentiationInterface/docs/Project.toml

@@ -9,6 +9,7 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

DifferentiationInterface/docs/src/api.md

@@ -68,7 +68,6 @@ jacobian
 jacobian!
 value_and_jacobian
 value_and_jacobian!
-MixedMode
 ```
 
 ## Second order
@@ -125,9 +124,10 @@ DifferentiationInterface.inner
 DifferentiateWith
 ```
 
-### Sparsity detection
+### Sparsity tools
 
 ```@docs
+MixedMode
 DenseSparsityDetector
 ```
 

DifferentiationInterface/docs/src/explanation/advanced.md

@@ -71,10 +71,26 @@ But after preparation, the more zeros are present in the matrix, the greater the
 ### Tuning the coloring algorithm
 
 The complexity of sparse Jacobians or Hessians grows with the number of distinct colors in a coloring of the sparsity pattern.
-To reduce this number of colors, [`GreedyColoringAlgorithm`](@ref) has two main settings: the order used for vertices and the decompression method.
+To reduce this number of colors, [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) has two main settings: the order used for vertices and the decompression method.
 Depending on your use case, you may want to modify either of these options to increase performance.
 See the documentation of [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl) for details.
 
+### Mixed mode
+
+When a Jacobian matrix has both dense rows and dense columns, it can be more efficient to use "mixed-mode" differentiation, a mixture of forward and reverse.
+The associated bidirectional coloring algorithm automatically decides how to cover the Jacobian using a set of columns (computed in forward mode) plus a set of rows (computed in reverse mode).
+This behavior is triggered as soon as you put a [`MixedMode`](@ref) object inside `AutoSparse`, like so:
+
+```julia
+AutoSparse(
+    MixedMode(forward_backend, reverse_backend);
+    sparsity_detector,
+    coloring_algorithm
+)
+```
+
+At the moment, mixed mode tends to work best when the [`GreedyColoringAlgorithm`](@extref SparseMatrixColorings.GreedyColoringAlgorithm) is provided with a [`RandomOrder`](@extref SparseMatrixColorings.RandomOrder) instead of the usual [`NaturalOrder`](@extref SparseMatrixColorings.NaturalOrder).
+
 ## Batch mode
 
 ### Multiple tangents

DifferentiationInterface/docs/src/tutorials/advanced.md

@@ -3,10 +3,12 @@
 We present contexts and sparsity handling with DifferentiationInterface.jl.
 
 ```@example tuto_advanced
+using ADTypes
 using BenchmarkTools
 using DifferentiationInterface
 import ForwardDiff, Zygote
-using SparseConnectivityTracer: TracerSparsityDetector
+using Random
+using SparseConnectivityTracer
 using SparseMatrixColorings
 ```
 
@@ -71,8 +73,8 @@ x = float.(1:8);
 ```
 
 ```@example tuto_advanced
-dense_first_order_backend = AutoForwardDiff()
-J_dense = jacobian(f_sparse_vector, dense_first_order_backend, x)
+dense_forward_backend = AutoForwardDiff()
+J_dense = jacobian(f_sparse_vector, dense_forward_backend, x)
 ```
 
 ```@example tuto_advanced
@@ -89,14 +91,14 @@ Recipe to create a sparse backend: combine a dense backend, a sparsity detector
 The following are reasonable defaults:
 
 ```@example tuto_advanced
-sparse_first_order_backend = AutoSparse(
-    dense_first_order_backend;
+sparse_forward_backend = AutoSparse(
+    dense_forward_backend;  # any object from ADTypes
     sparsity_detector=TracerSparsityDetector(),
     coloring_algorithm=GreedyColoringAlgorithm(),
 )
 
 sparse_second_order_backend = AutoSparse(
-    dense_second_order_backend;
+    dense_second_order_backend;  # any object from ADTypes or a SecondOrder from DI
     sparsity_detector=TracerSparsityDetector(),
     coloring_algorithm=GreedyColoringAlgorithm(),
 )
@@ -106,7 +108,7 @@ nothing  # hide
 Now the resulting matrices are sparse:
 
 ```@example tuto_advanced
-jacobian(f_sparse_vector, sparse_first_order_backend, x)
+jacobian(f_sparse_vector, sparse_forward_backend, x)
 ```
 
 ```@example tuto_advanced
@@ -123,7 +125,7 @@ Some result analysis functions from [SparseMatrixColorings.jl](https://github.co
 First, it records the sparsity pattern itself (the one returned by the detector).
 
 ```@example tuto_advanced
-jac_prep = prepare_jacobian(f_sparse_vector, sparse_first_order_backend, x)
+jac_prep = prepare_jacobian(f_sparse_vector, sparse_forward_backend, x)
 sparsity_pattern(jac_prep)
 ```
 
@@ -149,20 +151,20 @@ nothing  # hide
 ```
 
 ```@example tuto_advanced
-jac_prep_dense = prepare_jacobian(f_sparse_vector, dense_first_order_backend, zero(xbig))
-@benchmark jacobian($f_sparse_vector, $jac_prep_dense, $dense_first_order_backend, $xbig)
+jac_prep_dense = prepare_jacobian(f_sparse_vector, dense_forward_backend, zero(xbig))
+@benchmark jacobian($f_sparse_vector, $jac_prep_dense, $dense_forward_backend, $xbig)
 ```
 
 ```@example tuto_advanced
-jac_prep_sparse = prepare_jacobian(f_sparse_vector, sparse_first_order_backend, zero(xbig))
-@benchmark jacobian($f_sparse_vector, $jac_prep_sparse, $sparse_first_order_backend, $xbig)
+jac_prep_sparse = prepare_jacobian(f_sparse_vector, sparse_forward_backend, zero(xbig))
+@benchmark jacobian($f_sparse_vector, $jac_prep_sparse, $sparse_forward_backend, $xbig)
 ```
 
 Better memory use can be achieved by pre-allocating the matrix from the preparation result (so that it has the correct structure).
 
 ```@example tuto_advanced
 jac_buffer = similar(sparsity_pattern(jac_prep_sparse), eltype(xbig))
-@benchmark jacobian!($f_sparse_vector, $jac_buffer, $jac_prep_sparse, $sparse_first_order_backend, $xbig)
+@benchmark jacobian!($f_sparse_vector, $jac_buffer, $jac_prep_sparse, $sparse_forward_backend, $xbig)
 ```
 
 And for optimal speed, one should write non-allocating and type-stable functions.
@@ -184,7 +186,38 @@ ybig ≈ f_sparse_vector(xbig)
 In this case, the sparse Jacobian should also become non-allocating (for our specific choice of backend).
 
 ```@example tuto_advanced
-jac_prep_sparse_nonallocating = prepare_jacobian(f_sparse_vector!, zero(ybig), sparse_first_order_backend, zero(xbig))
+jac_prep_sparse_nonallocating = prepare_jacobian(f_sparse_vector!, zero(ybig), sparse_forward_backend, zero(xbig))
 jac_buffer = similar(sparsity_pattern(jac_prep_sparse_nonallocating), eltype(xbig))
-@benchmark jacobian!($f_sparse_vector!, $ybig, $jac_buffer, $jac_prep_sparse_nonallocating, $sparse_first_order_backend, $xbig)
+@benchmark jacobian!($f_sparse_vector!, $ybig, $jac_buffer, $jac_prep_sparse_nonallocating, $sparse_forward_backend, $xbig)
+```
+
+### Mixed mode
+
+Some Jacobians have a structure which includes dense rows and dense columns, like this one:
+
+```@example tuto_advanced
+arrowhead(x) = x .+ x[1] .+ vcat(sum(x), zeros(eltype(x), length(x)-1))
+
+jacobian_sparsity(arrowhead, x, TracerSparsityDetector())
+```
+
+In such cases, sparse AD is only beneficial in "mixed mode", where we combine a forward and a reverse backend.
+This is achieved using the [`MixedMode`](@ref) wrapper, for which we recommend a random coloring order (see [`RandomOrder`](@extref SparseMatrixColorings.RandomOrder)):
+
+```@example tuto_advanced
+sparse_mixed_backend = AutoSparse(
+    MixedMode(AutoForwardDiff(), AutoZygote()),
+    sparsity_detector=TracerSparsityDetector(),
+    coloring_algorithm=GreedyColoringAlgorithm(RandomOrder(MersenneTwister(), 0)),
+)
+```
+
+It unlocks a large speedup compared to pure forward mode, and the same would be true compared to reverse mode:
+
+```@example tuto_advanced
+@benchmark jacobian($arrowhead, prep, $sparse_forward_backend, $xbig) setup=(prep=prepare_jacobian(arrowhead, sparse_forward_backend, xbig))
+```
+
+```@example tuto_advanced
+@benchmark jacobian($arrowhead, prep, $sparse_mixed_backend, $xbig) setup=(prep=prepare_jacobian(arrowhead, sparse_mixed_backend, xbig))
 ```