is preparation thread safe?

[FerreolS]

Hi,
It is more a question than a issue as I didn't see any information in the doc:
Is the preparation thread safe?
This kind of code seems to work but is it legit for all possible backend?

julia> backendZ = AutoZygote()
AutoZygote()

julia> backendFD = AutoForwardDiff()
AutoForwardDiff()

julia> prepFD = prepare_gradient(sin, backendFD, x)
DifferentiationInterface.PullbackGradientPrep{Nothing, Float64, DifferentiationInterface.PushforwardPullbackPrep{Nothing, DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgPushforwardPrep{Nothing, ForwardDiff.Tag{typeof(sin), Float64}, ForwardDiff.Dual{ForwardDiff.Tag{typeof(sin), Float64}, Float64, 1}, Tuple{}}}}(Val{Nothing}(), 0.8414709848078965, DifferentiationInterface.PushforwardPullbackPrep{Nothing, DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgPushforwardPrep{Nothing, ForwardDiff.Tag{typeof(sin), Float64}, ForwardDiff.Dual{ForwardDiff.Tag{typeof(sin), Float64}, Float64, 1}, Tuple{}}}(Val{Nothing}(), DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgPushforwardPrep{Nothing, ForwardDiff.Tag{typeof(sin), Float64}, ForwardDiff.Dual{ForwardDiff.Tag{typeof(sin), Float64}, Float64, 1}, Tuple{}}(Val{Nothing}(), ForwardDiff.Tag{typeof(sin), Float64}, Dual{ForwardDiff.Tag{typeof(sin), Float64}}(1.0,1.0), ())))

julia> prepZ = prepare_gradient(sin, backendZ, x)
DifferentiationInterface.NoGradientPrep{Nothing}(Val{Nothing}())

julia> a = randn(1000);

julia> grad = similar(a);

julia> Threads.@threads for i in eachindex(a,grad); grad[i] = DifferentiationInterface.gradient(sin, prepZ, backendZ,a[i]) ; end

julia> Threads.@threads for i in eachindex(a,grad); grad[i] = DifferentiationInterface.gradient(sin, prepFD, backendFD,a[i]) ; end

[gdalle]

No, preparation is not thread-safe because the preparation object possibly gets mutated at every call. Actually, the code you wrote above probably doesn't give the correct answer inside grad for ForwardDiff. Here's a simpler example:

using Base.Threads
using DifferentiationInterface
import ForwardDiff

f(x) = sum(abs2, x)
backend = AutoForwardDiff()
prep = prepare_gradient(f, backend, zeros(2));

xs = [rand(2) for _ in 1:10_000];
gs = map(similar, xs)

@threads for i in eachindex(gs, xs)
    gradient!(f, gs[i], backend, xs[i])
end

all(map(==, 2 .* xs, gs))  # true

@threads for i in eachindex(gs, xs)
    gradient!(f, gs[i], prep, backend, xs[i])
end

all(map(==, 2 .* xs, gs))  # false

Would you like to open a PR adding this warning to the documentation? Here are two places where it could go:

• DifferentiationInterface.jl/DifferentiationInterface/src/docstrings.jl

  Lines 17 to 32 in ef77d13

  	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.
  	$(inplace ? "\nFor in-place functions, `y` is mutated by `f!` during preparation." : "")
  	!!! warning
  	The preparation result `prep` 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 becomes invalid and you need to run it again.
  	In some settings, invalid preparations may still give correct results (e.g. for backends that require no preparation), but this is not a semantic guarantee and should not be relied upon.
  	When `strict=Val(true)` (the default), type checking is enforced between preparation and execution (but size checking is left to the user).
  	While your code may work for different types by setting `strict=Val(false)`, this is not guaranteed by the API and can break without warning.
  	"""
  	end

• DifferentiationInterface.jl/DifferentiationInterface/docs/src/explanation/operators.md

  Line 121 in ef77d13

  ### Reusing preparation

[FerreolS]

That was what I presume.
I'll make a PR
Thanks