search

LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia
https://lux.csail.mit.edu/
MIT License
446 stars 50 forks source link

Support for inactive arguments in DifferentiationInterface #685

Closed gdalle closed 1 month ago

gdalle commented 1 month ago

https://github.com/LuxDL/Lux.jl/blob/44c7df2637a273a0615c85d325a53192d131398f/bench/helpers.jl#L56-L58

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

avik-pal commented 1 month ago

That comment was from @wsmoses

wsmoses commented 1 month ago

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 commented 1 month ago

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 commented 1 month ago

See https://github.com/gdalle/DifferentiationInterface.jl/issues/311 to continue the discussion