Support Enzyme.jl (#314)

* Enzyme WIP * cont. * cont. * Enzyme WIP * Add a sparse Jacobian with Enzyme.jl * Update tests * Update tests * Update nlpmodelstest.jl * Fix Project.toml * Fix the diff * Fix the tests * Fix the tests * Fix the tests * Fix tests * fix tests again * fix tests again * Add SparseEnzymeADHessian * Add SparseEnzymeADHessian * Add SparseEnzymeADHessian * Add SparseEnzymeADHessian * Add SparseEnzymeADHessian * nlp.f(x) -> f(x) * Fix Enzyme in sparse Hessian * Update the sparse Hessian * Fix the backends for Enzyme.jl * Update test/enzyme.jl * Update test/enzyme.jl * Update test/enzyme.jl * Update the function for the lagrangian * Update the function for the lagrangian * Update the function for the lagrangian * comment some tests * Uncomment some tests * Fix cx in Hessian * Fix * Fix the functor * Uncomment some code * fix test/Project.toml * Update the tests for Enzyme and Zygote * Update the documentation * Fix tests * Update nlp/basic.jl * FIx a typo * Fix documentation --------- Co-authored-by: Michel Schanen <[email protected]> Co-authored-by: Michel Schanen <[email protected]>
JuliaSmoothOptimizers · Dec 21, 2024 · 9988f4e · 9988f4e
1 parent 1547be0
commit 9988f4e
Show file tree

Hide file tree

Showing 21 changed files with 1,101 additions and 456 deletions.
diff --git a/docs/src/backend.md b/docs/src/backend.md
@@ -10,11 +10,11 @@ The functions used internally to define the NLPModel API and the possible backen
 | Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
 | --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
 | `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeADGradient` | -- |
-| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | -- | `SparseADJacobian` |
-| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | -- | `SparseADHessian`/`SparseReverseADHessian` |
-| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | -- | -- |
-| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | -- | -- |
-| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | -- | -- |
+| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `EnzymeADJacobian` | `SparseADJacobian` |
+| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | `EnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` |
+| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeADJprod` | -- |
+| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeADJtprod` | -- |
+| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | `EnzymeADHvprod` | -- |
 | `directional_second_derivative` | `ForwardDiffADGHjvprod` | -- | -- | -- | -- |
 
 The functions `hess_structure!`, `hess_coord!`, `jac_structure!` and `jac_coord!` defined in `ad.jl` are generic to all the backends for now.

diff --git a/docs/src/predefined.md b/docs/src/predefined.md
@@ -46,3 +46,12 @@ It is possible to use these pre-defined backends using the keyword argument `bac
 nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :optimized)
 get_adbackend(nlp)
 ```
+
+The backend `:enzyme` focuses on backend based on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+
+```@example ex1
+nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon, backend = :enzyme)
+get_adbackend(nlp)
+```
+
+The backend `:zygote` focuses on backend based on [Zygote.jl](https://github.com/FluxML/Zygote.jl).
diff --git a/docs/src/sparse.md b/docs/src/sparse.md
@@ -31,15 +31,15 @@ x = rand(T, 2)
 H = hess(nlp, x)
 ```
 
-The backends available for sparse derivatives (`SparseADJacobian`, `SparseADHessian`, and `SparseReverseADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
+The backends available for sparse derivatives (`SparseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
 These arguments specify the sparsity pattern detector and the coloring algorithm, respectively.
 
 - A **`detector`** must be of type `ADTypes.AbstractSparsityDetector`.
   The default detector is `TracerSparsityDetector()` from the package `SparseConnectivityTracer.jl`.
   Prior to version 0.8.0, the default was `SymbolicSparsityDetector()` from `Symbolics.jl`.
 
 - A **`coloring_algorithm`** must be of type `SparseMatrixColorings.GreedyColoringAlgorithm`.
-  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian`.
+  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`.
   These algorithms are provided by the package `SparseMatrixColorings.jl`.
 
 The `GreedyColoringAlgorithm{:direct}()` performs column coloring for Jacobians and star coloring for Hessians.