Multiple arguments / activities?

[gdalle]

At the moment the only viable way I see is to support one active and one inactive argument, and leave it to ComponentArrays to combine several arguments if necessary. It could be done by passing a tuple (x, c) instead of the input x. @wsmoses do you think that might get us 90% of the way to where we need to be?

See the following discussions:

• Enzyme hessian with closure: "You may be using a constant variable as temporary storage for active memory" #307
• Support for inactive arguments in DifferentiationInterface LuxDL/Lux.jl#685
• [WIP] Fresh attempt at DI integration SciML/OptimizationBase.jl#54
• Switch to DifferentiationInterface tpapp/LogDensityProblemsAD.jl#29

[wsmoses]

That would be an improvement but we need to make sure that the decomposition of the arguments doesn't result in a materialized intermediate variable

[gdalle]

Can you give an example and explain why that's a bad thing?

[wsmoses]

On movile so link chasing is hard but the logprovlemsAD has a thread where I described why closures were bad. The intermediate variable — whether tuple struct or closure — will have the same impact if not expanded

…

On Fri, Jun 7, 2024 at 12:40 AM Guillaume Dalle ***@***.***> wrote: Can you give an example and explain why that's a bad thing? — Reply to this email directly, view it on GitHub <#311 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXEDEIHL5QQ54Z6Z2SDZGDQMNAVCNFSM6AAAAABI5QNRGKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJTGUYTGNZYHA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[gdalle]

Yeah I added that link above, I guess I'm just wondering what you mean by intermediate variable in this setting. I don't see why grouping multiple arguments into one or unpacking them would require a closure

[jbrea]

and leave it to ComponentArrays to combine several arguments if necessary

Not sure, this would work for my use case. Typically, I have a function logp(data, model, parameters). data is constant and usually a vector of named tuples of whatever. model is usually a structure that contains configuration and state variables. The state variables need to be shadowed, as they change, but I don't care about their derivatives. The only thing I care about are the partial derivatives of logp with respect to parameters. Maybe it would make sense to follow Lux style and separate the model into two parts: constant configuration and non-contant state, which would result in logp(data, config, state, parameters). Would it really make sense to put data, model or data, config, state into a ComponentArray? How much would it limit the way data, config and state can be written?

To be clear: I don't care, if DifferentiationInterface covers my use case or not. I'm fine with using Enzyme directly. I was just exploring this direction and thought I report back my experience 🙂 .

[gdalle]

Thanks for your feedback! I imagine the following people may also have opinions:

• @avikpal for Lux
• @oxinabox for Zygote/Diffractor
• @willtebbutt for Tapir
• @adrhill forever <3

[gdalle]

As for your specific use case, we might have to make a distinction between "active / inactive" and "shadowed (mutated) / not shadowed"?

[willtebbutt]

No strongs views on my end -- activity analysis isn't presently really a thing in Tapir.jl, so happy to take the lead from others.

[gdalle]

Also note that the reason I'm mentioning ComponentArrays is because the original plan was never for DI to support an arbitrary number of arguments: that's what we agreed AbstractDifferentation would try to achieve based on DI (ping @mohamed82008). Hence I feel like disguising multiple arguments into one (x) or two (x and c for constant) is the best we can hope for in the short term.

[gdalle]

@Vaibhavdixit02 @ChrisRackauckas what kind of activity granularity would we need for Optimization.jl?

[ChrisRackauckas]

I think Optimization.jl is the one that should be reasonable to hit with the current interfaces. We can start there, and SciMLSensitivity.jl would need the full generality 😅

[wsmoses]

Forcing multiple arguments into a single struct — even two structs (one active one constant) will still result in the performance issues I outlined in the log density link.