Optimization.jl
Optimization.jl copied to clipboard
Support mixing manually defined functions with AD backends
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
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.
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
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.
This should work now https://github.com/SciML/Optimization.jl/blob/c62a2d71e80960fa6ea9e49ffbb4fa140e72a7c3/test/ADtests.jl#L251-L259 since SciML/Optimization.jl#342