Allow Duals and other types to pass through OptimizationFunction grad

[jClugstor]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Essentially checks if the input types are different from the prep object. If they are, just reprep for that call.

Specifically allows sensitivity analysis to use ForwardDiff with grad, to support SciML/SciMLSensitivity.jl#1444 .

[SebastianM-C]

out of curiosity, why is p also Dual? The optimization parameters can be pretty much any arbitrary user given struct.

[jClugstor]

The adjoint sensitivity finds the derivative of the solution of the optimization problem with respect to p. To find that we need:

$$\frac{\partial F}{\partial p}$$

where

$$F(w, p) = \begin{pmatrix} \nabla \mathcal{L}(w, p) \ g(w,p) \ h_\mathcal{I}(w,p) \end{pmatrix}^T$$

The problem I'm having is the $$\nabla \mathcal{L}(w, p)$$
because that's already a gradient with respect to $x$, and we need to then differentiate that with respect to $p$.

The optimization parameters can be pretty much any arbitrary user given struct.

Right, and if users want to get the sensitivity they'll either have to use types that are compatible with SciMLSensitivity or use SciMLStructures to tell SciMLSensitivity how to handle the user given struct. But that's already the case for ODEs and NonlinearSolves anyways, so it's still consistent with the rest of the ecosystem.

[jClugstor]

To be clear I'm not sure this PR is a good idea, and it's incomplete anyway since I would also need the constraint jacobians to be able to take Duals.

But this or another workaround is what would be needed if we want to be able to reuse the stuff that optimization builds inside of the adjoint.

[SebastianM-C]

ah, I see. It doesn't seem too invasive.

[jClugstor]

@SebastianM-C In order to support Enzyme for the optimization adjoint I needed to make some changes to the Enzyme extension. Mostly to allow the functions to accept other parameters, but there was also some boxing that Enzyme failed with, so I introduced some let blocks. There was also an issue where the closure was capturing a type value for the function annotation which Enzyme really did not like. So I changed it to precalculate so it would capture a Bool value instead. I don't think any of the changes are very intrusive but we definitely need to be careful with the Enzyme stuff, so if you see any issues let me know.

[jClugstor]

@ChrisRackauckas @SebastianM-C

In order for SciMLSensitivity to be able to reuse the closures built in Optimization, this PR has the changes needed.

One, make it so that they all can accept a p.

Two, make it so that they can accept Dual values. The way I do that in this PR is to use anyeltypedual to check if p is Dual at all, and if it is then we can't reuse the DI prep. At first I thought this would add some overhead but profiling and looking at the typed IR it turns out that if p is the same type as the prep the anyeltypedual gets compiled away so there's no runtime overhead.

The changes in the Enzyme extension are to make sure that the closures can take a p, and to prevent some Boxing from happening. Also there was an issue with nested Enzyme where the function annotation type was being captured that Enzyme couldn't handle so I fixed that as well.

[SebastianM-C]

One, make it so that they all can accept a p.

Ah, nice, I wanted to have that anyway.