Use a package extension for Enzyme.jl (#312)

Author: amontoison

.buildkite/pipeline.yml

@@ -23,16 +23,3 @@ steps:
       julia --color=yes --project=test -e 'using Pkg; Pkg.add("Enzyme"); Pkg.develop(path="."); Pkg.instantiate()'
       julia --color=yes --project=test -e 'include("test/enzyme.jl")'
     timeout_in_minutes: 30
-
-  - label: "CPUs -- Zygote.jl"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1.10"
-    agents:
-      queue: "juliaecosystem"
-      os: "linux"
-      arch: "x86_64"
-    command: |
-      julia --color=yes --project=test -e 'using Pkg; Pkg.add("Zygote"); Pkg.develop(path="."); Pkg.instantiate()'
-      julia --color=yes --project=test -e 'include("test/zygote.jl")'
-    timeout_in_minutes: 30

Project.toml

@@ -7,18 +7,25 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 
+[weakdeps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+
+[extensions]
+ADNLPModelsEnzymeExt = "Enzyme"
+
 [compat]
 ADTypes = "1.2.1"
-ForwardDiff = "0.9, 0.10, 1"
+ForwardDiff = "1"
 NLPModels = "0.21.5"
-Requires = "1"
 ReverseDiff = "1"
-SparseConnectivityTracer = "1.0"
+SparseConnectivityTracer = "1"
 SparseMatrixColorings = "0.4.21"
+Enzyme = "0.13.129"
+LinearAlgebra = "1.10"
+SparseArrays = "1.10"
 julia = "1.10"

docs/Project.toml

@@ -12,12 +12,11 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 SolverBenchmark = "581a75fa-a23a-52d0-a590-d6201de2218a"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 DataFrames = "1"
 Documenter = "1.0"
-ManualNLPModels = "0.1"
+ManualNLPModels = "0.2"
 NLPModels = "0.21.5"
 NLPModelsJuMP = "0.13"
 OptimizationProblems = "0.8"
@@ -26,4 +25,3 @@ Plots = "1"
 SolverBenchmark = "0.6"
 SparseConnectivityTracer = "1.0"
 SparseMatrixColorings = "0.4.21"
-Zygote = "0.6.62"

docs/src/backend.md

@@ -7,15 +7,15 @@ The backend information is in a structure [`ADNLPModels.ADModelBackend`](@ref) i
 
 The functions used internally to define the NLPModel API and the possible backends are defined in the following table:
 
-| Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
-| --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
-| `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeReverseADGradient` | -- |
-| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian` |
-| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | `SparseEnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` |
-| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeReverseADJprod` | -- |
-| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeReverseADJtprod` | -- |
-| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | `EnzymeReverseADHvprod` | -- |
-| `directional_second_derivative` | `ForwardDiffADGHjvprod` | -- | -- | -- | -- |
+| Functions | FowardDiff backends | ReverseDiff backends | Enzyme backend | Sparse backend |
+| --------- | ------------------- | -------------------- | -------------- | -------------- |
+| `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `EnzymeReverseADGradient` | -- |
+| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian` |
+| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `SparseEnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` |
+| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `EnzymeReverseADJprod` | -- |
+| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `EnzymeReverseADJtprod` | -- |
+| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | `EnzymeReverseADHvprod` | -- |
+| `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.
 

docs/src/index.md

@@ -24,8 +24,8 @@ ADNLPModels.jl functionalities are extended by other packages that are not autom
 In other words, you sometimes need to load the desired package separately to access some functionalities.
 
 ```julia
-using ADNLPModels # load only the default functionalities
-using Zygote # load the Zygote backends
+using ADNLPModels  # load only the default functionalities
+using Enzyme       # load the Enzyme backends
 ```
 
 Versions compatibility for the extensions are available in the file `test/Project.toml`.

docs/src/predefined.md

@@ -56,5 +56,3 @@ get_adbackend(nlp)
 
 !!! danger
     The interface for Enzyme.jl is still under development.
-
-The backend `:zygote` focuses on backend based on [Zygote.jl](https://github.com/FluxML/Zygote.jl).