Support for inactive arguments in DifferentiationInterface

[gdalle]

Lux.jl/bench/helpers.jl

Lines 56 to 58 in 44c7df2

	# - Note DI's present design requires the creation of closures and prevents
	# marking some arguments as inactive. As such, in its current form it is
	# ill-suited as a high-performance/high-compatibility interface for Enzyme.

Do you think supporting one active and one inactive argument in DI would be enough?

[avik-pal]

That comment was from @wsmoses

[wsmoses]

I think the bigger issue is the closure.

So the question is whether DI can construct the same multi-arg multi-activity Enzyme autodiff call, without constructing a closure.

[gdalle]

Generic multi-arg is much much much harder to test, and so we'll probably never support it in DI (although AbstractDifferentiation may pick it up based on DI, ping @mohamed82008). Luckily, ComponentArrays takes some of that pain away by allowing us to group arguments in many cases.

I'm not gonna have an activity hierarchy as detailed as that of Enzyme, but I was thinking that allowing one active and one inactive argument might be enough for most practical purposes. We can group active arguments in one ComponentVector and inactive arguments in another if there are several. And the syntax could be very similar to the current DI syntax, with e.g.

jacobian(f, backend, x)  # one active argument
jacobian(f, backend, (x, c))  # one active argument and one constant

[gdalle]

See JuliaDiff/DifferentiationInterface.jl#311 to continue the discussion