Possible support for Dual types for Automatic Differentiation

[albertomercurio]

Hello,

I was trying to use this package together with the ForwardDiff.jl package, in order to use automatic differentiation, but it failed because the input types are restricted to just floats or integers but not Duals, which are the base types for automatic differentiation.

I was wondering if relaxing the input types would be possible, in order to add the support to automatic differentiation, which is very used in the machine learning field but also when fitting some functions. Indeed, I was trying to use this package to get the smallest eigenvalues of a matrix as a function of some parameter, and the fitting these parameters with some experimental data’s. This fit procedure would be much faster with automatic differentiation.

[amontoison]

@albertomercurio We did a package DiffKrylov.jl with @michel2323 to support AD with Krylov.jl.
We used another strategy than relaxing the type to keep the performance on all architectures (CPU / GPU).

https://github.com/JuliaSmoothOptimizers/DiffKrylov.jl

[albertomercurio]

Oh, good to know.

Why don’t you make an extension to this package, instead of creating a new package for these purposes?

[amontoison]

It's still a work-in-progress project and we worked on a paper for the conference AD 2024. It was easier for the review to not include it in Krylov.jl yet.

[michel2323]

Let us know if you need any help with a specific setup. DiffKrylov should have support for ForwardDiff, though we mostly used Enzyme and reverse mode. In general, it is better not to differentiate through an iterative solver even in forward mode. Your stopping criteria will be based on the primal value, and your duals might be inaccurate/not converged.

It was also not added as an extension right away because Krylov.jl is very lightweight in its dependencies. Eventually, DiffKrylov.jl would only provide rules for ChainRules.jl and EnzymeRules.jl.

A hint to the KrylovPreconditioners.jl (also WIP) package in case you want a fast, portable, and differentiable preconditioner setup.

[albertomercurio]

Ok thanks for the clarification. I have already noticed KrylovPreconditioners.jl, but I still never used it. I will do it for sure in the future.

BTW, just for my clarification, making it as an extension will not influence the dependencies of the package, right? It will continue to be a lightweight package, no? This is just for clarifying my concept of extension.

[michel2323]

That is true. However, there are, AFAIK, two things to keep in mind with extensions.

1. In extensions, you can only add methods to a function. You cannot define new functions.
2. If common packages like ForwardDiff trigger the extension, it might trigger the extension's loading (and lots of dependencies), although in the code, you are using Krylov.jl and ForwardDiff.jl independently.

[amontoison]

The other drawback of an extension is that it's only related to one package. For example, if you need both ForwardDiff.jl and ChainRules.jl to implement a new method of cg that works with AD types, you can't use an extension.

[albertomercurio]

Thank you, very clear.