docs: improve and automate docstrings (#742)

Author: gdalle

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -36,6 +36,7 @@ using ADTypes:
 using LinearAlgebra: dot
 
 include("compat.jl")
+include("docstrings.jl")
 
 include("first_order/mixed_mode.jl")
 include("second_order/second_order.jl")
@@ -45,7 +46,7 @@ include("utils/traits.jl")
 include("utils/basis.jl")
 include("utils/batchsize.jl")
 include("utils/check.jl")
-include("utils/printing.jl")
+include("utils/errors.jl")
 include("utils/context.jl")
 include("utils/linalg.jl")
 include("utils/sparse.jl")

DifferentiationInterface/src/docstrings.jl

@@ -0,0 +1,51 @@
+function docstring_preptype(preptype::AbstractString, operator::AbstractString)
+    return """
+        $(preptype)
+
+    Abstract type for additional information needed by [`$(operator)`](@ref) and its variants.
+    """
+end
+
+function samepoint_warning(samepoint::Bool)
+    if samepoint
+        ", _if they are applied at the same point `x` and with the same `contexts`_"
+    else
+        ""
+    end
+end
+
+function docstring_prepare(operator; samepoint=false, inplace=false)
+    return """
+    Create a `prep` object that can be given to [`$(operator)`](@ref) and its variants to speed them up$(samepoint_warning(samepoint)).
+
+    Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.
+
+    !!! warning
+        The preparation result is only reusable as long as the arguments to `$operator` do not change type or size, and the function and backend themselves are not modified.
+        Otherwise, preparation will be invalidated and you will need to run it again.
+    $(inplace ? "\nFor in-place functions, `y` is mutated by `f!` during preparation." : "")
+    """
+end
+
+function docstring_prepare!(operator)
+    return """
+    Same behavior as [`prepare_$(operator)`](@ref) but can resize the contents of an existing `prep` object to avoid some allocations.
+
+    There is no guarantee that `prep` will be mutated, or that performance will be improved compared to preparation from scratch.
+
+    !!! danger
+        Compared to when `prep` was first created, the only authorized modification is a size change for input `x` or output `y`.
+        Any other modification (like a change of type for the input) is not supported and will give erroneous results.
+
+    !!! danger
+        For efficiency, this function needs to rely on backend package internals, therefore it not protected by semantic versioning.
+    """
+end
+
+function docstring_preparation_hint(operator::AbstractString; same_point=false)
+    if same_point
+        return "To improve performance via operator preparation, refer to [`prepare_$(operator)`](@ref) and [`prepare_$(operator)_same_point`](@ref)."
+    else
+        return "To improve performance via operator preparation, refer to [`prepare_$(operator)`](@ref)."
+    end
+end

DifferentiationInterface/src/first_order/derivative.jl

@@ -4,24 +4,15 @@
     prepare_derivative(f,     backend, x, [contexts...]) -> prep
     prepare_derivative(f!, y, backend, x, [contexts...]) -> prep
 
-Create a `prep` object that can be given to [`derivative`](@ref) and its variants.
-
-!!! warning
-    If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.
-    For in-place functions, `y` is mutated by `f!` during preparation.
+$(docstring_prepare("derivative"; inplace=true))
 """
 function prepare_derivative end
 
 """
     prepare!_derivative(f,     prep, backend, x, [contexts...]) -> new_prep
     prepare!_derivative(f!, y, prep, backend, x, [contexts...]) -> new_prep
 
-Same behavior as [`prepare_derivative`](@ref) but can modify an existing `prep` object to avoid some allocations.
-
-There is no guarantee that `prep` will be mutated, or that performance will be improved compared to preparation from scratch.
-
-!!! danger
-    For efficiency, this function needs to rely on backend package internals, therefore it not protected by semantic versioning.
+$(docstring_prepare!("derivative"))
 """
 function prepare!_derivative end
 
@@ -31,7 +22,7 @@ function prepare!_derivative end
 
 Compute the value and the derivative of the function `f` at point `x`.
 
-$(document_preparation("derivative"))
+$(docstring_preparation_hint("derivative"))
 """
 function value_and_derivative end
 
@@ -41,7 +32,7 @@ function value_and_derivative end
 
 Compute the value and the derivative of the function `f` at point `x`, overwriting `der`.
 
-$(document_preparation("derivative"))
+$(docstring_preparation_hint("derivative"))
 """
 function value_and_derivative! end
 
@@ -51,7 +42,7 @@ function value_and_derivative! end
 
 Compute the derivative of the function `f` at point `x`.
 
-$(document_preparation("derivative"))
+$(docstring_preparation_hint("derivative"))
 """
 function derivative end
 
@@ -61,7 +52,7 @@ function derivative end
 
 Compute the derivative of the function `f` at point `x`, overwriting `der`.
 
-$(document_preparation("derivative"))
+$(docstring_preparation_hint("derivative"))
 """
 function derivative! end
 

DifferentiationInterface/src/first_order/gradient.jl

@@ -3,22 +3,14 @@
 """
     prepare_gradient(f, backend, x, [contexts...]) -> prep
 
-Create a `prep` object that can be given to [`gradient`](@ref) and its variants.
-
-!!! warning
-    If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.
+$(docstring_prepare("gradient"))
 """
 function prepare_gradient end
 
 """
     prepare!_gradient(f, prep, backend, x, [contexts...]) -> new_prep
 
-Same behavior as [`prepare_gradient`](@ref) but can modify an existing `prep` object to avoid some allocations.
-
-There is no guarantee that `prep` will be mutated, or that performance will be improved compared to preparation from scratch.
-
-!!! danger
-    For efficiency, this function needs to rely on backend package internals, therefore it not protected by semantic versioning.
+$(docstring_prepare!("gradient"))
 """
 function prepare!_gradient end
 
@@ -27,7 +19,7 @@ function prepare!_gradient end
 
 Compute the value and the gradient of the function `f` at point `x`.
 
-$(document_preparation("gradient"))
+$(docstring_preparation_hint("gradient"))
 """
 function value_and_gradient end
 
@@ -36,7 +28,7 @@ function value_and_gradient end
 
 Compute the value and the gradient of the function `f` at point `x`, overwriting `grad`.
 
-$(document_preparation("gradient"))
+$(docstring_preparation_hint("gradient"))
 """
 function value_and_gradient! end
 
@@ -45,7 +37,7 @@ function value_and_gradient! end
 
 Compute the gradient of the function `f` at point `x`.
 
-$(document_preparation("gradient"))
+$(docstring_preparation_hint("gradient"))
 """
 function gradient end
 
@@ -54,7 +46,7 @@ function gradient end
 
 Compute the gradient of the function `f` at point `x`, overwriting `grad`.
 
-$(document_preparation("gradient"))
+$(docstring_preparation_hint("gradient"))
 """
 function gradient! end
 

DifferentiationInterface/src/first_order/jacobian.jl

@@ -4,24 +4,15 @@
     prepare_jacobian(f,     backend, x, [contexts...]) -> prep
     prepare_jacobian(f!, y, backend, x, [contexts...]) -> prep
 
-Create a `prep` object that can be given to [`jacobian`](@ref) and its variants.
-
-!!! warning
-    If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.
-    For in-place functions, `y` is mutated by `f!` during preparation.
+$(docstring_prepare("jacobian"; inplace=true))
 """
 function prepare_jacobian end
 
 """
     prepare!_jacobian(f,     prep, backend, x, [contexts...]) -> new_prep
     prepare!_jacobian(f!, y, prep, backend, x, [contexts...]) -> new_prep
 
-Same behavior as [`prepare_jacobian`](@ref) but can modify an existing `prep` object to avoid some allocations.
-
-There is no guarantee that `prep` will be mutated, or that performance will be improved compared to preparation from scratch.
-
-!!! danger
-    For efficiency, this function needs to rely on backend package internals, therefore it not protected by semantic versioning.
+$(docstring_prepare!("jacobian"))
 """
 function prepare!_jacobian end
 
@@ -31,7 +22,7 @@ function prepare!_jacobian end
 
 Compute the value and the Jacobian matrix of the function `f` at point `x`.
 
-$(document_preparation("jacobian"))
+$(docstring_preparation_hint("jacobian"))
 """
 function value_and_jacobian end
 
@@ -41,7 +32,7 @@ function value_and_jacobian end
 
 Compute the value and the Jacobian matrix of the function `f` at point `x`, overwriting `jac`.
     
-$(document_preparation("jacobian"))
+$(docstring_preparation_hint("jacobian"))
 """
 function value_and_jacobian! end
 
@@ -51,7 +42,7 @@ function value_and_jacobian! end
 
 Compute the Jacobian matrix of the function `f` at point `x`.
 
-$(document_preparation("jacobian"))
+$(docstring_preparation_hint("jacobian"))
 """
 function jacobian end
 
@@ -61,7 +52,7 @@ function jacobian end
 
 Compute the Jacobian matrix of the function `f` at point `x`, overwriting `jac`.
 
-$(document_preparation("jacobian"))
+$(docstring_preparation_hint("jacobian"))
 """
 function jacobian! end
 

DifferentiationInterface/src/first_order/mixed_mode.jl

@@ -1,10 +1,10 @@
 """
     MixedMode
 
-Combination of a forward and a reverse mode backend for mixed-mode Jacobian computation.
+Combination of a forward and a reverse mode backend for mixed-mode sparse Jacobian computation.
 
 !!! danger
-    `MixedMode` backends only support [`jacobian`](@ref) and its variants.
+    `MixedMode` backends only support [`jacobian`](@ref) and its variants, and it should be used inside an [`AutoSparse`](@extref ADTypes.AutoSparse) wrapper.
 
 # Constructor
 
@@ -20,8 +20,24 @@ struct MixedMode{F<:AbstractADType,R<:AbstractADType} <: AbstractADType
     end
 end
 
+"""
+    forward_backend(m::MixedMode)
+
+Return the forward-mode part of a `MixedMode` backend.
+"""
 forward_backend(m::MixedMode) = m.forward
+
+"""
+    reverse_backend(m::MixedMode)
+
+Return the reverse-mode part of a `MixedMode` backend.
+"""
 reverse_backend(m::MixedMode) = m.reverse
 
+"""
+    ForwardAndReverseMode <: ADTypes.AbstractMode
+
+Appropriate mode type for `MixedMode` backends.
+"""
 struct ForwardAndReverseMode <: ADTypes.AbstractMode end
 ADTypes.mode(::MixedMode) = ForwardAndReverseMode()

DifferentiationInterface/src/first_order/pullback.jl

@@ -4,36 +4,23 @@
     prepare_pullback(f,     backend, x, ty, [contexts...]) -> prep
     prepare_pullback(f!, y, backend, x, ty, [contexts...]) -> prep
 
-Create a `prep` object that can be given to [`pullback`](@ref) and its variants.
-
-!!! warning
-    If the function changes in any way, the result of preparation will be invalidated, and you will need to run it again.
-    For in-place functions, `y` is mutated by `f!` during preparation.
+$(docstring_prepare("pullback"; inplace=true))
 """
 function prepare_pullback end
 
 """
     prepare!_pullback(f,     prep, backend, x, ty, [contexts...]) -> new_prep
     prepare!_pullback(f!, y, prep, backend, x, ty, [contexts...]) -> new_prep
 
-Same behavior as [`prepare_pullback`](@ref) but can modify an existing `prep` object to avoid some allocations.
-
-There is no guarantee that `prep` will be mutated, or that performance will be improved compared to preparation from scratch.
-
-!!! danger
-    For efficiency, this function needs to rely on backend package internals, therefore it not protected by semantic versioning.
+$(docstring_prepare!("pullback"))
 """
 function prepare!_pullback end
 
 """
     prepare_pullback_same_point(f,     backend, x, ty, [contexts...]) -> prep_same
     prepare_pullback_same_point(f!, y, backend, x, ty, [contexts...]) -> prep_same
 
-Create an `prep_same` object that can be given to [`pullback`](@ref) and its variants _if they are applied at the same point `x` and with the same `contexts`_.
-
-!!! warning
-    If the function or the point changes in any way, the result of preparation will be invalidated, and you will need to run it again.
-    For in-place functions, `y` is mutated by `f!` during preparation.
+$(docstring_prepare("pullback"; samepoint=true, inplace=true))
 """
 function prepare_pullback_same_point end
 
@@ -43,7 +30,7 @@ function prepare_pullback_same_point end
 
 Compute the value and the pullback of the function `f` at point `x` with a tuple of tangents `ty`.
 
-$(document_preparation("pullback"; same_point=true))
+$(docstring_preparation_hint("pullback"; same_point=true))
 
 !!! tip 
     Pullbacks are also commonly called vector-Jacobian products or VJPs.
@@ -60,7 +47,7 @@ function value_and_pullback end
 
 Compute the value and the pullback of the function `f` at point `x` with a tuple of tangents `ty`, overwriting `dx`.
 
-$(document_preparation("pullback"; same_point=true))
+$(docstring_preparation_hint("pullback"; same_point=true))
 
 !!! tip 
     Pullbacks are also commonly called vector-Jacobian products or VJPs.
@@ -74,7 +61,7 @@ function value_and_pullback! end
 
 Compute the pullback of the function `f` at point `x` with a tuple of tangents `ty`.
 
-$(document_preparation("pullback"; same_point=true))
+$(docstring_preparation_hint("pullback"; same_point=true))
 
 !!! tip 
     Pullbacks are also commonly called vector-Jacobian products or VJPs.
@@ -88,7 +75,7 @@ function pullback end
 
 Compute the pullback of the function `f` at point `x` with a tuple of tangents `ty`, overwriting `dx`.
 
-$(document_preparation("pullback"; same_point=true))
+$(docstring_preparation_hint("pullback"; same_point=true))
 
 !!! tip 
     Pullbacks are also commonly called vector-Jacobian products or VJPs.