Optimization.jl icon indicating copy to clipboard operation
Optimization.jl copied to clipboard

Support mixing manually defined functions with AD backends

Open gaurav-arya opened this issue 1 year ago • 3 comments

The instantiate function for most AD backends only populates a field if it is nothing. On the other hand, the instantiate function for NoAD only processes those fields that are not nothing.

It seems like these two pieces of logic should work together rather than be disjoint. In particular, it could be nice to allow the user to define some fields and leave the rest to the AD backend, but this doesn't really work in the current design because the processing of those fields (which currently involves creating closures over p it seems) is inside instantiate_function for the NoAD backend which does not get run when an AD backend is set.

A possible design could be: rename definitions by backends to __instantiate_function, and have an instantiate_function implementation that always does the first part and then calls out to the AD backend for the second part. The __instantiate_function for NoAD would just do nothing then. Perhaps the "first path" could also include changing out-of-place user functions into in-place versions, see https://github.com/SciML/Optimization.jl/issues/312

gaurav-arya avatar Jul 06 '22 17:07 gaurav-arya

In particular, it could be nice to allow the user to define some fields and leave the rest to the AD backend, but this doesn't really work in the current design because the processing of those fields (which currently involves creating closures over p it seems) is inside instantiate_function for the NoAD backend which does not get run when an AD backend is set.

Indeed, this is an issue and your fix is a good idea for how to fix it.

ChrisRackauckas avatar Jul 06 '22 17:07 ChrisRackauckas

I agree, this approach sounds good, I have also been thinking about this and SciML/Optimization.jl#191 - we should extend the thought to allow passing the choice of AD backend for constraint, I was thinking each field can take the Auto* objects but I feel that might be overkill so instead we can add a kwarg like cons_adtype that can be used to pass the different choice

Vaibhavdixit02 avatar Jul 07 '22 04:07 Vaibhavdixit02

I wonder whether allowing choice per function is overkill. It could have a pretty big compile time effect because then you get combinatoric explosion in the amount of type combinations and thus code combinations allowed. I'm not sure how that would play out in practice, but I'd almost say it would be best to section that off to a AutoGeneralAD() or something so that it's an opt-in special case.

Anyways, I think we might as well get the current setup solid, and then think about generalizing after constraints etc. are fully supported. We have a lot of things to do just to get the current scope solid.

ChrisRackauckas avatar Jul 07 '22 21:07 ChrisRackauckas

This should work now https://github.com/SciML/Optimization.jl/blob/c62a2d71e80960fa6ea9e49ffbb4fa140e72a7c3/test/ADtests.jl#L251-L259 since SciML/Optimization.jl#342

Vaibhavdixit02 avatar Aug 25 '22 11:08 Vaibhavdixit02