Remove Zygote from everything

docs/src/ext/ahmc.md

@@ -15,7 +15,7 @@ chain = sample(post, NUTS(0.8), 10_000; n_adapts=5_000)
 
 In addition our sample call has a few additional keyword arguments:
 
- - `adtype = Val(:Zygote)`: The autodiff package to use. Currently the default is `Zygote` and we recommend using this. Note that you must load Zygote before calling `sample`.
+ - `adtype = Val(:Enzyme)`: The autodiff package to use. Currently the only options is `Enzyme`. Note that you must load Enzyme before calling `sample`.
  - `saveto = MemoryStore()`: Specifies how to store the samples. The default is `MemoryStore` which stores the samples directly in RAM. For large models this is not a good idea. To save samples periodically to disk use [`DiskStore`](@ref), and then load the results with `load_samples`.
 
 Note that like most `AbstractMCMC` samplers the initial location can be specified with the `initial_params` argument.
@@ -26,7 +26,7 @@ Note that like most `AbstractMCMC` samplers the initial location can be specifie
 ```julia
 using Comrade
 using AdvancedHMC
-using Zygote
+using Enzyme
 
 # Some stuff to create a posterior object
 post # of type Comrade.Posterior

docs/src/ext/optimization.md

@@ -13,9 +13,10 @@ To see what optimizers are available and what options are available, please see
 using Comrade
 using Optimization
 using OptimizationOptimJL
+using Enzyme
 
 # Some stuff to create a posterior object
 post # of type Comrade.Posterior
 
-xopt, sol = comrade_opt(post, LBFGS(); adtype=Val(:Zygote))
+xopt, sol = comrade_opt(post, LBFGS(); adtype=Val(:Enzyme))
 ```

docs/src/introduction.md

@@ -8,12 +8,16 @@ Comrade is a Bayesian differentiable modular modeling framework for use with ver
 The goal is to allow the user to easily combine and modify a set of primitive models
 to construct complicated source structures. The benefit of this approach is that it is straightforward to construct different source models out of these primitives. Namely, an end-user does
 not have to create a separate source "model" every time they
-change the model specification. Additionally, most models currently implemented are differentiable with at Zygote and sometimes ForwardDiff[^2]. This allows for gradient accelerated optimization and sampling (e.g., HMC) to be used with little
-effort by the end user. To sample from the posterior, we provide a somewhat barebones interface since, most of the time, and we don't require the additional features offered by most PPLs. Additionally, the overhead introduced by PPLs tends to be rather large. In the future, we may revisit this as
-Julia's PPL ecosystem matures.
+change the model specification. Additionally, most models currently implemented are differentiable with at Enzyme. This allows for gradient accelerated optimization and sampling (e.g., HMC) to be used with little
+effort by the end user. 
 
-!!! note
-    The primitives the Comrade defines, however, would allow for it to be easily included in PPLs like [`Turing`](https://github.com/TuringLang/Turing.jl).
+!!! warn
+    As of 0.11 Comrade is only tested with Enzyme. ForwardDiff may work for some
+    models, but it is not tested and there are known problems with e.g., closures
+    due to the use of a sparse cholesky factorization. Additionally essentially no
+    models will work with Zygote due to the heavy use of mutation in the codebase.
+    The use of mutation is necessary for performance reasons, and we have no plans
+    to change this in the future.
 
 
 ## Tutorials
@@ -28,7 +32,8 @@ the image and instrumental effects like gain ratios and leakage terms in [Polari
 
 This repository has recently moved to [ColPrac](https://github.com/SciML/ColPrac). If you would like to contribute please feel free to open a issue or pull-request.
 
-[^2]: As of 0.9 Comrade switched to using full covariance closures. As a result this requires a sparse cholesky solve in the likelihood evaluation which requires which isn't defined in ForwardDiff. As a result we recommend using Zygote which does work and often is similarly performant (reverse pass is 3-6x slower compared to the forward pass).
+
+```@meta
 
 ## Requirements
 

examples/advanced/HybridImaging/main.jl

@@ -178,8 +178,8 @@ fig |> DisplayAs.PNG |> DisplayAs.Text #hide
 # To use this we use the [`comrade_opt`](@ref) function
 using Optimization
 using OptimizationOptimJL
-using Zygote
-xopt, sol = comrade_opt(post, LBFGS(), Optimization.AutoZygote(); initial_params=prior_sample(rng, post), maxiters=1000, g_tol=1e0)
+using Enzyme
+xopt, sol = comrade_opt(post, LBFGS(), AutoEnzyme(Enzyme.Reverse); initial_params=prior_sample(rng, post), maxiters=1000, g_tol=1e0)
 
 
 # First we will evaluate our fit by plotting the residuals

examples/intermediate/ClosureImaging/main.jl

@@ -147,11 +147,11 @@ post = VLBIPosterior(skym, dlcamp, dcphase)
 # To optimize our posterior `Comrade` provides the `comrade_opt` function. To use this
 # functionality a user first needs to import `Optimization.jl` and the optimizer of choice.
 # In this tutorial we will use Optim.jl's L-BFGS optimizer, which is defined in the sub-package
-# OptimizationOptimJL. We also need to import Zygote to allow for automatic differentiation.
+# OptimizationOptimJL. We also need to import Enzyme to allow for automatic differentiation.
 using Optimization
 using OptimizationOptimJL
-using Zygote
-xopt, sol = comrade_opt(post, LBFGS(), Optimization.AutoZygote(); maxiters=1000)
+using Enzyme
+xopt, sol = comrade_opt(post, LBFGS(), AutoEnzyme(Enzyme.Reverse); maxiters=1000)
 
 
 # First we will evaluate our fit by plotting the residuals

examples/intermediate/PolarizedImaging/main.jl

@@ -320,9 +320,7 @@ using BenchmarkTools
 
 # Now we optimize. Unlike other imaging examples, we move straight to gradient optimizers
 # due to the higher dimension of the space. In addition the only AD package that can currently
-# work with the polarized Comrade posterior is Zygote. Note that in the future we expect to
-# shift entirely to Enzyme, and in fact large portions of Comrade's AD already uses Enzyme
-# through custom rules.
+# work with the polarized Comrade posterior is Enzyme.
 using Optimization
 using OptimizationOptimisers
 using Enzyme

ext/ComradeAdvancedHMCExt.jl

@@ -36,15 +36,15 @@ end
 
 function AbstractMCMC.Sample(
             rng::Random.AbstractRNG, tpost::Comrade.TransformedVLBIPosterior,
-            sampler::AbstractHMCSampler; adtype=Val(:Zygote), initial_params=nothiing, kwargs...)
+            sampler::AbstractHMCSampler; adtype=Val(:Enzyme), initial_params=nothiing, kwargs...)
     ∇ℓ = ADgradient(adtype, tpost)
     θ0 = initialize_params(tpost, initial_params)
     model, smplr = make_sampler(rng, ∇ℓ, sampler, θ0)
     return AbstractMCMC.Sample(rng, model, smplr; initial_params=θ0, kwargs...)
 end
 
 """
-    sample(rng, post::VLBIPosterior, sampler::AbstractHMCSampler, nsamples, args...;saveto=MemoryStore(), adtype=Val(:Zygote), initial_params=nothing, kwargs...)
+    sample(rng, post::VLBIPosterior, sampler::AbstractHMCSampler, nsamples, args...;saveto=MemoryStore(), adtype=Val(:Enzyme), initial_params=nothing, kwargs...)
 
 Sample from the posterior `post` using the sampler `sampler` for `nsamples` samples. Additional
 arguments are forwarded to AbstractMCMC.sample. If `saveto` is a DiskStore, the samples will be
@@ -59,15 +59,15 @@ saved to disk. If `initial_params` is not `nothing` then the sampler will start
 ## Keyword Arguments
 
  - `saveto`: If a DiskStore, the samples will be saved to disk, if [`MemoryStore`](@ref) the samples will be stored in memory/ram.
- - `adtype`: The automatic differentiation type to use. The default if Zygote which is the recommended choice for Comrade currently.
+ - `adtype`: The automatic differentiation type to use. The default if Enzyme which is the recommended choice for Comrade currently.
  - `initial_params`: The initial parameters to start the sampler from. If `nothing` then the sampler will start from a random point in the prior.
  - `kwargs`: Additional keyword arguments to pass to the sampler. Examples include `n_adapts` which is the total number of samples to use for adaptation.
     To see the others see the AdvancedHMC documentation.
 """
 function AbstractMCMC.sample(
         rng::Random.AbstractRNG, post::Comrade.VLBIPosterior,
         sampler::AbstractHMCSampler, nsamples, args...;
-        saveto=MemoryStore(), adtype=Val(:Zygote), initial_params=nothing, kwargs...)
+        saveto=MemoryStore(), adtype=Val(:Enzyme), initial_params=nothing, kwargs...)
 
     saveto isa DiskStore && return sample_to_disk(rng, post, sampler, nsamples, args...; outdir=saveto.name, output_stride=min(saveto.stride, nsamples), adtype, initial_params, kwargs...)
 
@@ -90,7 +90,7 @@ end
 function initialize(rng::Random.AbstractRNG, tpost::Comrade.TransformedVLBIPosterior,
     sampler::AbstractHMCSampler, nsamples, outbase, args...;
     n_adapts = min(nsamples÷2, 1000),
-    adtype = Val(:Zygote),
+    adtype = Val(:Enzyme),
     initial_params=nothing, outdir = "Results",
     output_stride=min(100, nsamples),
     restart = false,
@@ -158,7 +158,7 @@ end
 
 function sample_to_disk(rng::Random.AbstractRNG, post::Comrade.VLBIPosterior,
                         sampler::AbstractHMCSampler, nsamples, args...;
-                        adtype = Val(:Zygote),
+                        adtype = Val(:Enzyme),
                         n_adapts = min(nsamples÷2, 1000),
                         initial_params=nothing, outdir = "Results",
                         restart=false,

ext/ComradeOptimizationExt.jl

@@ -46,7 +46,7 @@ Optimize the posterior `post` using the `opt` optimizer.
  - `opt` : The optimizer to use. This can be any optimizer from `Optimization.jl`.
  - `adtype` : The automatic differentiation type to use. The default is `nothing` which means
     no automatic differentiation is used. To specify to use automatic differentiation
-    set `adtype`. For example if you wish to use `Zygote` set `adtype=Optimization.AutoZygote()`.
+    set `adtype`. For example if you wish to use `Enzyme` set `adtype=Optimization.AutoEnzyme(Enzyme.Reverse)`.
  - `args` : Additional arguments passed to the `Optimization`, `solve` method
 
 ## Keyword Arguments

ext/ComradePigeonsExt.jl

@@ -33,7 +33,7 @@ Pigeons.default_explorer(::Comrade.TransformedVLBIPosterior{P,<:HypercubeTransfo
     SliceSampler()
 
 Pigeons.default_explorer(::Comrade.TransformedVLBIPosterior{P,<:TransformVariables.AbstractTransform}) where {P} =
-    Pigeons.AutoMALA(;default_autodiff_backend = :Zygote)
+    Pigeons.AutoMALA(;default_autodiff_backend = :Enzyme)
 
 function Pigeons.default_reference(tpost::Comrade.TransformedVLBIPosterior)
     t = tpost.transform

lib/ComradeAHMC/test/runtests.jl

@@ -1,5 +1,5 @@
 using Pyehtim, Comrade, ComradeAHMC, Distributions, VLBIImagePriors
-using Zygote
+using Enzyme
 using Test
 
 include(joinpath(@__DIR__, "../../../test/test_util.jl"))
@@ -21,8 +21,8 @@ include(joinpath(@__DIR__, "../../../test/test_util.jl"))
           ξ2 = -0.5376269092893298,
           x = 1.451956089157719e-10,
           y = 1.455983181049137e-10),)
-    s1 = AHMC(autodiff=Val(:Zygote))
-    s2 = AHMC(autodiff=Val(:Zygote))
+    s1 = AHMC(autodiff=Val(:Enzyme))
+    s2 = AHMC(autodiff=Val(:Enzyme))
     s3 = AHMC()
     hchain = sample(post, s1, 1_000; n_adapts=500, progress=false)
     hchain = sample(post, s1, 1_000; n_adapts=500, progress=false, initial_params=x0)