Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support Enzyme.jl #314

Merged
merged 44 commits into from
Dec 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
44 commits
Select commit Hold shift + click to select a range
b615dd8
Enzyme WIP
michel2323 Nov 26, 2024
4a186be
cont.
michel2323 Nov 26, 2024
52a79bd
cont.
michel2323 Nov 26, 2024
5de2577
Enzyme WIP
michel2323 Nov 27, 2024
aee2282
Add a sparse Jacobian with Enzyme.jl
amontoison Nov 27, 2024
2e08e9a
Update tests
amontoison Nov 27, 2024
298e218
Update tests
amontoison Nov 27, 2024
2cc2372
Update nlpmodelstest.jl
amontoison Nov 27, 2024
f768b4e
Fix Project.toml
amontoison Nov 27, 2024
668ce53
Fix the diff
amontoison Nov 27, 2024
293a879
Fix the tests
amontoison Nov 27, 2024
875f250
Fix the tests
amontoison Nov 28, 2024
9865f1a
Fix the tests
amontoison Nov 28, 2024
22e5dd7
Fix tests
amontoison Nov 30, 2024
9b63fb6
fix tests again
amontoison Nov 30, 2024
c86d848
fix tests again
amontoison Nov 30, 2024
659c5b3
Add SparseEnzymeADHessian
amontoison Dec 2, 2024
2b3e803
Add SparseEnzymeADHessian
amontoison Dec 2, 2024
2bfb90c
Add SparseEnzymeADHessian
amontoison Dec 2, 2024
2795795
Add SparseEnzymeADHessian
amontoison Dec 2, 2024
1900a0d
Add SparseEnzymeADHessian
amontoison Dec 2, 2024
a95f515
nlp.f(x) -> f(x)
amontoison Dec 2, 2024
9ccd2f4
Fix Enzyme in sparse Hessian
michel2323 Dec 2, 2024
e1c7747
Update the sparse Hessian
amontoison Dec 2, 2024
b98429b
Fix the backends for Enzyme.jl
amontoison Dec 2, 2024
a813360
Update test/enzyme.jl
amontoison Dec 2, 2024
e8a2d0c
Update test/enzyme.jl
amontoison Dec 2, 2024
d11ce2f
Update test/enzyme.jl
amontoison Dec 2, 2024
e750118
Update the function for the lagrangian
amontoison Dec 3, 2024
09b4793
Update the function for the lagrangian
amontoison Dec 3, 2024
8ca8044
Update the function for the lagrangian
amontoison Dec 3, 2024
9f1c2cc
comment some tests
amontoison Dec 3, 2024
fed8fa3
Uncomment some tests
amontoison Dec 3, 2024
11ea3eb
Fix cx in Hessian
michel2323 Dec 3, 2024
0fd47a2
Fix
michel2323 Dec 3, 2024
5915c81
Fix the functor
amontoison Dec 4, 2024
0355f5f
Uncomment some code
amontoison Dec 20, 2024
f870f70
fix test/Project.toml
amontoison Dec 20, 2024
921ad6c
Update the tests for Enzyme and Zygote
amontoison Dec 21, 2024
d62b308
Update the documentation
amontoison Dec 21, 2024
b4e88ab
Fix tests
amontoison Dec 21, 2024
7eb8c16
Update nlp/basic.jl
amontoison Dec 21, 2024
3f1ff12
FIx a typo
amontoison Dec 21, 2024
ed926ed
Fix documentation
amontoison Dec 21, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions docs/src/backend.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
9 changes: 9 additions & 0 deletions docs/src/predefined.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).
4 changes: 2 additions & 2 deletions docs/src/sparse.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading
Loading