Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion DifferentiationInterface/src/DifferentiationInterface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand Down
51 changes: 51 additions & 0 deletions DifferentiationInterface/src/docstrings.jl
Original file line number Diff line number Diff line change
@@ -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`_"

Check warning on line 11 in DifferentiationInterface/src/docstrings.jl

View check run for this annotation

Codecov / codecov/patch

DifferentiationInterface/src/docstrings.jl#L11

Added line #L11 was not covered by tests
else
""

Check warning on line 13 in DifferentiationInterface/src/docstrings.jl

View check run for this annotation

Codecov / codecov/patch

DifferentiationInterface/src/docstrings.jl#L13

Added line #L13 was not covered by tests
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
21 changes: 6 additions & 15 deletions DifferentiationInterface/src/first_order/derivative.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand Down
20 changes: 6 additions & 14 deletions DifferentiationInterface/src/first_order/gradient.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand Down
21 changes: 6 additions & 15 deletions DifferentiationInterface/src/first_order/jacobian.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand Down
20 changes: 18 additions & 2 deletions DifferentiationInterface/src/first_order/mixed_mode.jl
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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()
27 changes: 7 additions & 20 deletions DifferentiationInterface/src/first_order/pullback.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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.
Expand Down
Loading
Loading