DifferentiationInterface do-block best practices

Specific Domains Numerics
,
1

Is it better / recommended to use do-blocks that capture variables or should I use the Context mechanism provided by DifferentiationInterface?

Even though this

pack = (y0 = 1.0, y1 = 2.0)
x1 = [2.0, 4.0]
jac = DifferentiationInterface.jacobian(<some_backend>, x1) do x
  return (x[1] - pack.y0)^2 + (x[2]-pack.y1)^2
end

computes the same thing as this

pack = (y0 = 1.0, y1 = 2.0)
x1 = [2.0, 4.0]
jac = DifferentiationInterface.jacobian(<some_backend>, x1, Constant(pack.y0), Constant(pack.y1) do x, y0, y1
  return (x[1] - y0)^2 + (x[2]-y1)^2
end

is one preferable / better than the other? I suppose the second is much more amenable to use with the preparation mechanism.

2

Contexts are in general preferable because

  • they allow stateless functions, which is an implicit assumption of many backends and makes code cleaner
  • for some backends, captured variables would need to be tweaked to accommodate differentiation (eg DI.Cache with ForwardDiff.jl to work with dual numbers)
  • for some backends, context annotations enable performance optimizations (eg DI.Constant with Enzyme.jl)
2 Likes