API reference
DifferentiationInterfaceTest
— ModuleDifferentiationInterfaceTest
Testing and benchmarking utilities for automatic differentiation in Julia.
Exports
Entry points
DifferentiationInterfaceTest.Scenario
— TypeScenario{op,pl_op,pl_fun}
Store a testing scenario composed of a function and its input + output for a given operator.
This generic type should never be used directly: use the specific constructor corresponding to the operator you want to test, or a predefined list of scenarios.
Type parameters
op
: one of:pushforward
,:pullback
,:derivative
,:gradient
,:jacobian
,:second_derivative
,:hvp
,:hessian
pl_op
: either:in
(forop!(f, result, backend, x)
) or:out
(forresult = op(f, backend, x)
)pl_fun
: either:in
(forf!(y, x)
) or:out
(fory = f(x)
)
Constructors
Scenario{op,pl_op}(f, x; tang, contexts, res1, res2)
Scenario{op,pl_op}(f!, y, x; tang, contexts, res1, res2)
Fields
f::Any
: functionf
(ifargs==1
) orf!
(ifargs==2
) to applyx::Any
: primal inputy::Any
: primal outputtang::Union{Nothing, NTuple{N, T} where {N, T}}
: tangents for pushforward, pullback or HVPcontexts::Tuple
: contexts (if applicable)res1::Any
: first-order result of the operator (if applicable)res2::Any
: second-order result of the operator (if applicable)
DifferentiationInterfaceTest.test_differentiation
— Functiontest_differentiation(
backends::Vector{<:ADTypes.AbstractADType};
...
) -> Union{Nothing, DataFrames.DataFrame}
test_differentiation(
backends::Vector{<:ADTypes.AbstractADType},
scenarios::Vector{<:Scenario};
correctness,
type_stability,
allocations,
benchmark,
excluded,
detailed,
logging,
isapprox,
atol,
rtol,
scenario_intact,
sparsity,
ignored_modules,
function_filter,
skip_allocations,
count_calls,
benchmark_test
) -> Union{Nothing, DataFrames.DataFrame}
Apply a list of backends
on a list of scenarios
, running a variety of different tests and/or benchmarks.
Return
This function always creates and runs a @testset
, though its contents may vary.
- if
benchmark == :none
, it returnsnothing
. - if
benchmark != :none
, it returns aDataFrame
of benchmark results, whose columns correspond to the fields ofDifferentiationBenchmarkDataRow
.
Positional arguments
backends::Vector{<:AbstractADType}
: the backends to testscenarios::Vector{<:Scenario}
: the scenarios on which to test them (defaults to the output ofdefault_scenarios()
)
Keyword arguments
Test categories:
correctness=true
: whether to compare the differentiation results with the theoretical values specified in each scenariotype_stability=:none
: whether (and how) to check type stability of operators with JET.jl.allocations=:none
: whether (and how) to check allocations inside operators with AllocCheck.jlbenchmark=:none
: whether (and how) to benchmark operators with Chairmarks.jl
For type_stability
, allocations
and benchmark
, the possible values are :none
, :prepared
or :full
. Each setting tests/benchmarks a different subset of calls:
kwarg | prepared operator | unprepared operator | preparation |
---|---|---|---|
:none | no | no | no |
:prepared | yes | no | no |
:full | yes | yes | yes |
Misc options:
excluded::Vector{Symbol}
: list of operators to exclude, such asFIRST_ORDER
orSECOND_ORDER
detailed=false
: whether to create a detailed or condensed testsetlogging=false
: whether to log progress
Correctness options:
isapprox=isapprox
: function used to compare objects approximately, with the standard signatureisapprox(x, y; atol, rtol)
atol=0
: absolute precision for correctness testing (when comparing to the reference outputs)rtol=1e-3
: relative precision for correctness testing (when comparing to the reference outputs)scenario_intact=true
: whether to check that the scenario remains unchanged after the operators are appliedsparsity=false
: whether to check sparsity patterns for Jacobians / Hessians
Type stability options:
ignored_modules=nothing
: list of modules that JET.jl should ignorefunction_filter
: filter for functions that JET.jl should ignore (with a reasonable default)
Benchmark options:
count_calls=true
: whether to also count function calls during benchmarkingbenchmark_test=true
: whether to include tests which succeed iff benchmark doesn't error
test_differentiation(
backend::ADTypes.AbstractADType,
args...;
kwargs...
) -> Union{Nothing, DataFrames.DataFrame}
Shortcut for a single backend.
DifferentiationInterfaceTest.benchmark_differentiation
— Functionbenchmark_differentiation(
backends,
scenarios::Vector{<:Scenario};
benchmark,
excluded,
logging,
count_calls,
benchmark_test
) -> Union{Nothing, DataFrames.DataFrame}
Shortcut for test_differentiation
with only benchmarks and no correctness or type stability checks.
Specifying the set of scenarios is mandatory for this function.
DifferentiationInterfaceTest.FIRST_ORDER
— ConstantFIRST_ORDER = [:pushforward, :pullback, :derivative, :gradient, :jacobian]
List of all first-order operators, to facilitate exclusion during tests.
DifferentiationInterfaceTest.SECOND_ORDER
— ConstantSECOND_ORDER = [:hvp, :second_derivative, :hessian]
List of all second-order operators, to facilitate exclusion during tests.
Utilities
DifferentiationInterfaceTest.DifferentiationBenchmarkDataRow
— TypeDifferentiationBenchmarkDataRow
Ad-hoc storage type for differentiation benchmarking results.
Fields
backend::ADTypes.AbstractADType
: backend used for benchmarkingscenario::Scenario
: scenario used for benchmarkingoperator::Symbol
: differentiation operator used for benchmarking, e.g.:gradient
or:hessian
prepared::Union{Nothing, Bool}
: whether the operator had been preparedcalls::Int64
: number of calls to the differentiated function for one call to the operatorsamples::Int64
: number of benchmarking samples takenevals::Int64
: number of evaluations used for averaging in each sampletime::Float64
: minimum runtime over all samples, in secondsallocs::Float64
: minimum number of allocations over all samplesbytes::Float64
: minimum memory allocated over all samples, in bytesgc_fraction::Float64
: minimum fraction of time spent in garbage collection over all samples, between 0.0 and 1.0compile_fraction::Float64
: minimum fraction of time spent compiling over all samples, between 0.0 and 1.0
See the documentation of Chairmarks.jl for more details on the measurement fields.
Pre-made scenario lists
The precise contents of the scenario lists are not part of the API, only their existence.
DifferentiationInterfaceTest.default_scenarios
— Functiondefault_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with standard array types.
DifferentiationInterfaceTest.sparse_scenarios
— Functionsparse_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with sparse array types, focused on sparse Jacobians and Hessians.
DifferentiationInterfaceTest.component_scenarios
— Functioncomponent_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with component array types from ComponentArrays.jl.
This function requires ComponentArrays.jl to be loaded (it is implemented in a package extension).
DifferentiationInterfaceTest.gpu_scenarios
— Functiongpu_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with GPU array types from JLArrays.jl.
This function requires JLArrays.jl to be loaded (it is implemented in a package extension).
DifferentiationInterfaceTest.static_scenarios
— Functionstatic_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with static array types from StaticArrays.jl.
This function requires StaticArrays.jl to be loaded (it is implemented in a package extension).
Internals
This is not part of the public API.
Base.zero
— Methodzero(scen::Scenario)
Return a new Scenario
identical to scen
except for the first- and second-order results which are set to zero.
DifferentiationInterfaceTest.allocfree_scenarios
— Functionallocfree_scenarios(rng::AbstractRNG=default_rng())
Create a vector of Scenario
s with functions that do not allocate.
At the moment, second-order scenarios are excluded.
DifferentiationInterfaceTest.batchify
— Methodbatchify(scen::Scenario)
Return a new Scenario
identical to scen
except for the tangents tang
and associated results res1
/ res2
, which are duplicated (batch mode).
Only works if scen
is a pushforward
, pullback
or hvp
scenario.
DifferentiationInterfaceTest.cachify
— Methodcachify(scen::Scenario)
Return a new Scenario
identical to scen
except for the function f
, which is made to accept an additional cache argument a
to store the result before it is returned.
DifferentiationInterfaceTest.change_function
— Methodchange_function(scen::Scenario, new_f)
Return a new Scenario
identical to scen
except for the function f
which is changed to new_f
.
DifferentiationInterfaceTest.closurify
— Methodclosurify(scen::Scenario)
Return a new Scenario
identical to scen
except for the function f
which is made to close over differentiable data.
DifferentiationInterfaceTest.constantify
— Methodconstantify(scen::Scenario)
Return a new Scenario
identical to scen
except for the function f
, which is made to accept an additional constant argument a
by which the output is multiplied. The output and result fields are updated accordingly.
DifferentiationInterfaceTest.flux_isapprox
— Functionflux_isapprox(x, y; atol, rtol)
Approximate comparison function to use in correctness tests with gradients of Flux.jl networks.
DifferentiationInterfaceTest.flux_scenarios
— Functionflux_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with neural networks from Flux.jl.
This function requires FiniteDifferences.jl and Flux.jl to be loaded (it is implemented in a package extension).
These scenarios are still experimental and not part of the public API. Their ground truth values are computed with finite differences, and thus subject to imprecision.
DifferentiationInterfaceTest.lux_isapprox
— Functionlux_isapprox(x, y; atol, rtol)
Approximate comparison function to use in correctness tests with gradients of Lux.jl networks.
DifferentiationInterfaceTest.lux_scenarios
— Functionlux_scenarios(rng=Random.default_rng())
Create a vector of Scenario
s with neural networks from Lux.jl.
This function requires ComponentArrays.jl, ForwardDiff.jl, Lux.jl and LuxTestUtils.jl to be loaded (it is implemented in a package extension).
These scenarios are still experimental and not part of the public API.