API reference

DifferentiationInterfaceTest

Testing and benchmarking utilities for automatic differentiation in Julia.

Scenario{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 (for op!(f, result, backend, x)) or :out (for result = op(f, backend, x))
• pl_fun: either :in (for f!(y, x)) or :out (for y = f(x))

Constructors

Scenario{op,pl_op}(
    f, x, [t], contexts...;
    prep_args, res1, res2, name
)

Scenario{op,pl_op}(
    f!, y, x, [t,] contexts...;
    prep_args, res1, res2, name
)

Default values:

• prep_args = the result of zero applied to each execution argument
• res1 = res2 = nothing
• name = nothing

Fields

• f::Any: function f (if pl_fun==:out) or f! (if pl_fun==:in) to apply

• y::Any: primal output

• x::Any: primal input

• t::Union{Nothing, NTuple{N, T} where {N, T}}: tangents (if applicable)

• contexts::Tuple: contexts (if applicable)

• res1::Any: first-order result of the operator (if applicable)

• res2::Any: second-order result of the operator (if applicable)

• prep_args::NamedTuple: named tuple of arguments passed to preparation, without the function - the required keys are a subset of (; y, x, t, contexts) depending on the operator

• name::Union{Nothing, String}: name of the scenario for display in test sets and dataframes

compute_results(scen::Scenario, backend::AbstractADType)

Return a scenario identical to scen but where the first- and second-order results res1 and res2 have been computed with the given differentiation backend.

Useful for comparison of outputs between backends.

test_differentiation(
    backends::Vector{<:ADTypes.AbstractADType};
    ...
) -> Union{Nothing, DataFrames.DataFrame}
test_differentiation(
    backends::Vector{<:ADTypes.AbstractADType},
    scenarios::Vector{<:Scenario};
    testset_name,
    correctness,
    type_stability,
    allocations,
    benchmark,
    excluded,
    detailed,
    logging,
    isapprox,
    atol,
    rtol,
    scenario_intact,
    sparsity,
    reprepare,
    ignored_modules,
    function_filter,
    skip_allocations,
    count_calls,
    benchmark_test,
    benchmark_seconds,
    benchmark_aggregation,
    adaptive_batchsize
) -> 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 returns nothing.
• if benchmark != :none, it returns a DataFrame of benchmark results, whose columns correspond to the fields of DifferentiationBenchmarkDataRow.

Positional arguments

• backends::Vector{<:AbstractADType}: the backends to test
• scenarios::Vector{<:Scenario}: the scenarios on which to test these backends. Defaults to a standard set of first- and second-order scenarios, whose contents are not part of the public API and may change without notice.

Keyword arguments

• testset_name=nothing: how to display the test set

Test categories:

• correctness=true: whether to compare the differentiation results with the theoretical values specified in each scenario
• type_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.jl
• benchmark=: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:

Misc options:

• excluded::Vector{Symbol}: list of operators to exclude, such as FIRST_ORDER or SECOND_ORDER
• detailed=false: whether to create a detailed or condensed testset
• logging=false: whether to log progress

Correctness options:

• isapprox=isapprox: function used to compare objects approximately, with the standard signature isapprox(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 applied
• sparsity=false: whether to check sparsity patterns for Jacobians / Hessians
• reprepare::Bool=true: whether to modify preparation before testing when the preparation arguments have the wrong size

Type stability options:

• ignored_modules=nothing: list of modules that JET.jl should ignore
• function_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 benchmarking
• benchmark_test=true: whether to include tests which succeed iff benchmark doesn't error
• benchmark_seconds=1: how long to run each benchmark for
• benchmark_aggregation=minimum: function used to aggregate sample measurements

Batch size options

• adaptive_batchsize=true: whether to cap the backend's preset batch size (when it exists) to prevent errors on small inputs

test_differentiation(
    backend::ADTypes.AbstractADType,
    args...;
    kwargs...
) -> Union{Nothing, DataFrames.DataFrame}

Shortcut for a single backend.

benchmark_differentiation(
    backends,
    scenarios::Vector{<:Scenario};
    testset_name,
    benchmark,
    excluded,
    logging,
    count_calls,
    benchmark_test,
    benchmark_seconds,
    benchmark_aggregation,
    adaptive_batchsize
) -> 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.

FIRST_ORDER = [:pushforward, :pullback, :derivative, :gradient, :jacobian]

List of all first-order operators, to facilitate exclusion during tests.

SECOND_ORDER = [:hvp, :second_derivative, :hessian]

List of all second-order operators, to facilitate exclusion during tests.

DifferentiationBenchmarkDataRow

Ad-hoc storage type for differentiation benchmarking results.

Fields

• backend::ADTypes.AbstractADType: backend used for benchmarking

• scenario::Scenario: scenario used for benchmarking

• operator::Symbol: differentiation operator used for benchmarking, e.g. :gradient or :hessian

• prepared::Union{Nothing, Bool}: whether the operator had been prepared

• calls::Int64: number of calls to the differentiated function for one call to the operator

• samples::Int64: number of benchmarking samples taken

• evals::Int64: number of evaluations used for averaging in each sample

• time::Any: aggregated runtime over all samples, in seconds

• allocs::Any: aggregated number of allocations over all samples

• bytes::Any: aggregated memory allocated over all samples, in bytes

• gc_fraction::Any: aggregated fraction of time spent in garbage collection over all samples, between 0.0 and 1.0

• compile_fraction::Any: aggregated 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.

zero(scen::Scenario)

Return a new Scenario identical to scen except for the first- and second-order results which are set to zero.

allocfree_scenarios()

Create a vector of Scenarios with functions that do not allocate.

batchify(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.

cachify(scen::Scenario)

Return a new Scenario identical to scen except for the function f, which is made to accept an additional cache argument to store the result before it is returned.

If tup=true the cache is a tuple of arrays, otherwise just an array.

change_function(scen::Scenario, new_f)

Return a new Scenario identical to scen except for the function f which is changed to new_f.

closurify(scen::Scenario)

Return a new Scenario identical to scen except for the function f which is made to close over differentiable data.

complex_scenarios()

Create a vector of first-order Scenarios with complex-valued array types and holomorphic functions only.

complex_sparse_scenarios()

Create a vector of Jacobian Scenarios with complex-valued array types and holomorphic functions only.

component_scenarios()

Create a vector of Scenarios with component array types from ComponentArrays.jl.

constantify(scen::Scenario)

Return a new Scenario identical to scen except for the function f, which is made to accept an additional constant argument by which the output is multiplied. The output and result fields are updated accordingly.

constantorcachify(scen::Scenario)

Return a new Scenario identical to scen except for the function f, which is made to accept an additional "constant or cache" argument.

default_scenarios()

Create a vector of Scenarios with standard array types.

flux_isapprox(x, y; atol, rtol)

Approximate comparison function to use in correctness tests with gradients of Flux.jl networks.

flux_scenarios(rng=Random.default_rng())

Create a vector of Scenarios with neural networks from Flux.jl.

gpu_scenarios()

Create a vector of Scenarios with GPU array types from JLArrays.jl.

lux_isapprox(x, y; atol, rtol)

Approximate comparison function to use in correctness tests with gradients of Lux.jl networks.

lux_scenarios(rng=Random.default_rng())

Create a vector of Scenarios with neural networks from Lux.jl.

sparse_scenarios()

Create a vector of Scenarios with sparse array types, focused on sparse Jacobians and Hessians.

static_scenarios()

Create a vector of Scenarios with static array types from StaticArrays.jl.