diff --git a/.github/workflows/AutoDiffIntegration.yml b/.github/workflows/AutoDiffIntegration.yml index 7a2071596..f745bf3a4 100644 --- a/.github/workflows/AutoDiffIntegration.yml +++ b/.github/workflows/AutoDiffIntegration.yml @@ -5,6 +5,8 @@ on: AD: required: true type: string +permissions: + contents: read concurrency: # Skip intermediate builds: always. @@ -14,18 +16,14 @@ concurrency: jobs: test: - name: Julia ${{ matrix.AD }} - ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }} - runs-on: ${{ matrix.os }} + name: Julia ${{ matrix.AD }} - ${{ matrix.version }} - ubuntu-latest - ${{ github.event_name }} + runs-on: ubuntu-latest strategy: fail-fast: false matrix: version: - 'lts' - '1.11' - os: - - ubuntu-latest - - macOS-latest - - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 diff --git a/.github/workflows/DynamicPPL.yml b/.github/workflows/DynamicPPL.yml index 459ec2afd..a48a3cec4 100644 --- a/.github/workflows/DynamicPPL.yml +++ b/.github/workflows/DynamicPPL.yml @@ -6,6 +6,8 @@ on: tags: ['*'] pull_request: workflow_dispatch: +permissions: + contents: read concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. @@ -13,18 +15,14 @@ concurrency: cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} jobs: test: - name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }} - runs-on: ${{ matrix.os }} + name: Julia ${{ matrix.version }} - ubuntu-latest - ${{ github.event_name }} + runs-on: ubuntu-latest strategy: fail-fast: false matrix: version: - 'lts' - '1.11' - os: - - ubuntu-latest - - macOS-latest - - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 3368a965a..77296d139 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -6,6 +6,8 @@ on: tags: ['*'] pull_request: workflow_dispatch: +permissions: + contents: read concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. @@ -13,18 +15,14 @@ concurrency: cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} jobs: test: - name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }} - runs-on: ${{ matrix.os }} + name: Julia ${{ matrix.version }} - ubuntu-latest - ${{ github.event_name }} + runs-on: ubuntu-latest strategy: fail-fast: false matrix: version: - 'lts' - '1.11' - os: - - ubuntu-latest - - macOS-latest - - windows-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 diff --git a/.github/workflows/Zygote.yml b/.github/workflows/Zygote.yml deleted file mode 100644 index 6ed2aa31c..000000000 --- a/.github/workflows/Zygote.yml +++ /dev/null @@ -1,14 +0,0 @@ -name: Zygote Integration Tests -on: - push: - branches: - - main - tags: ['*'] - pull_request: - workflow_dispatch: - -jobs: - test: - uses: ./.github/workflows/AutoDiffIntegration.yml - with: - AD: Zygote diff --git a/README.md b/README.md index fd1da9caa..4a5214f39 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,6 @@ |:---------------------------------------------------------- |:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) | [![ForwardDiff](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/ForwardDiff.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/ForwardDiff.yml?query=branch%3Amain) | | [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl) | [![ReverseDiff](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/ReverseDiff.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/ReverseDiff.yml?query=branch%3Amain) | -| [Zygote](https://github.com/FluxML/Zygote.jl) | [![Zygote](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Zygote.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Zygote.yml?query=branch%3Amain) | | [Mooncake](https://github.com/chalk-lab/Mooncake.jl) | [![Mooncake](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Mooncake.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Mooncake.yml?query=branch%3Amain) | | [Enzyme](https://github.com/EnzymeAD/Enzyme.jl) | [![Enzyme](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Enzyme.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/Enzyme.yml?query=branch%3Amain) | diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 1d07fd975..4044f6732 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -35,7 +35,7 @@ Evaluate the value and gradient of a function `f` at `x` using the automatic dif # Arguments - `ad::ADTypes.AbstractADType`: automatic differentiation backend. Currently supports - `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, + `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, `ADTypes.AutoMooncake()` and `ADTypes.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), diff --git a/src/algorithms/repgradelbo.jl b/src/algorithms/repgradelbo.jl index ee2fccb8a..29a20cb0b 100644 --- a/src/algorithms/repgradelbo.jl +++ b/src/algorithms/repgradelbo.jl @@ -13,7 +13,7 @@ Evidence lower-bound objective with the reparameterization gradient formulation[ # Requirements - The variational approximation ``q_{\\lambda}`` implements `rand`. - The target distribution and the variational approximation have the same support. -- The target `LogDensityProblem` should satisfy either of the following: The target has a capability of at least `LogDensityProblems.LogDensityOrder{1}()` and the AD backend is one of `ReverseDiff`, `Zygote`, `Mooncake`, and `AutoEnzyme` in reverse mode so that `ADTypes.mode(adtype) == ADTypes.ReverseMode` is true. (In this case, `AdvancedVI` will take advantage of the existing `LogDensityProblems.logdensity_and_gradient`.) Otherwise, `LogDensityProblems.logdensity` should be differentiable under the selected AD backend. +- The target `LogDensityProblem` should satisfy either of the following: The target has a capability of at least `LogDensityProblems.LogDensityOrder{1}()` and the AD backend is one of `ReverseDiff`, `Mooncake`, and `AutoEnzyme` in reverse mode so that `ADTypes.mode(adtype) == ADTypes.ReverseMode` is true. (In this case, `AdvancedVI` will take advantage of the existing `LogDensityProblems.logdensity_and_gradient`.) Otherwise, `LogDensityProblems.logdensity` should be differentiable under the selected AD backend. - The sampling process `rand(q)` must be differentiable by the selected AD backend. Depending on the options, additional requirements on ``q_{\\lambda}`` may apply. @@ -53,10 +53,10 @@ function init( prob else if !( - adtype isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoMooncake,<:AutoEnzyme} && + adtype isa Union{<:AutoReverseDiff,<:AutoMooncake,<:AutoEnzyme} && ADTypes.mode(adtype) isa ADTypes.ReverseMode ) - @info "The capability of the supplied target `LogDensityProblem` $(capability) is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoZygote`, `AutoMooncake`, or `AutoEnzyme` in reverse mode." + @info "The capability of the supplied target `LogDensityProblem` $(capability) is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode." end MixedADLogDensityProblem(prob) end diff --git a/test/Project.toml b/test/Project.toml index ce8ea455b..356cbf4e1 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,7 +21,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ADTypes = "0.2.1, 1" @@ -46,5 +45,4 @@ Statistics = "1" StatsBase = "0.34" Test = "1" Tracker = "0.2.20" -Zygote = "0.6.63, 0.7" julia = "1.10, 1.11.2" diff --git a/test/general/mixedad_logdensity.jl b/test/general/mixedad_logdensity.jl index bee48a397..af6d65401 100644 --- a/test/general/mixedad_logdensity.jl +++ b/test/general/mixedad_logdensity.jl @@ -31,8 +31,8 @@ function mixedad_test_fwd(x, prob) )/2 end -# MixedADLogDensityProblem only supports ReverseDiff, Zygote, Enzyme, Mooncake in reverse-mode -if (AD isa Union{<:AutoReverseDiff,<:AutoZygote,<:AutoEnzyme,<:AutoMooncake}) && +# MixedADLogDensityProblem only supports ReverseDiff, Enzyme, Mooncake in reverse-mode +if (AD isa Union{<:AutoReverseDiff,<:AutoEnzyme,<:AutoMooncake}) && (ADTypes.mode(AD) isa ADTypes.ReverseMode) @testset "MixedADLogDensityProblem" begin model = MixedADTestModel() diff --git a/test/runtests.jl b/test/runtests.jl index e1e230eb9..b50e69bfc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -27,9 +27,6 @@ const AD = if AD_str == "ReverseDiff" AutoReverseDiff() elseif AD_str == "ForwardDiff" AutoForwardDiff() -elseif AD_str == "Zygote" - using Zygote - AutoZygote() elseif AD_str == "Mooncake" using Mooncake AutoMooncake(; config=Mooncake.Config()) @@ -39,6 +36,8 @@ elseif AD_str == "Enzyme" mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=Enzyme.Const, ) +else + throw(ArgumentError("Unsupported AD backend for tests: $AD_str")) end if GROUP == "DynamicPPL"