Remove Zygote dependency and update documentation for ForwardDiff integration #393

AoifeHughes · 2025-06-16T10:57:05Z

Remove Zygote Support from `Bijectors.jl` - Addresses: #374

Summary

This PR removes Zygote as a supported automatic differentiation backend from Bijectors.jl, aligning with the policy that only ForwardDiff, ReverseDiff, and Mooncake are officially supported for DynamicPPL.

Changes Made

Files Deleted

ext/BijectorsZygoteExt.jl
Complete extension module (199 lines) containing 20+ custom @adjoint implementations for Zygote.

Project Configuration

Project.toml
Removed Zygote from:
- [weakdeps] section
- [extensions] section (BijectorsZygoteExt = "Zygote")
- [compat] constraints
- [extras] dependencies
test/Project.toml
- Removed Zygote dependency and compatibility constraints
docs/Project.toml
- Removed Zygote dependency and compatibility constraints

Source Code

test/runtests.jl
- Removed using Zygote import statement
test/ad/utils.jl
- Removed all Zygote testing logic from test_ad():
  - Removed :Zygote from broken backends list
  - Removed Zygote gradient testing code blocks
src/bijectors/pd.jl
- Updated comment to remove Zygote reference
src/chainrules.jl
- Generalized comment about AD issues with @debug macro

Documentation

docs/src/examples.md
- Replaced Zygote.gradient example with ForwardDiff.gradient in the normalizing flows training example
CLAUDE.md
- Updated to reflect currently supported AD backends

Impact

Removed: ~200 lines of Zygote-specific automatic differentiation code
Files modified: 10 files across project configuration, source code, tests, and documentation
Maintained support: ForwardDiff, ReverseDiff, and Mooncake remain fully supported
Breaking change: Users relying on Zygote will need to migrate to supported AD backends

Testing

All Zygote-specific tests have been removed
Remaining AD tests continue to validate ForwardDiff, ReverseDiff, and Mooncake functionality
No functional changes to the core bijector interface

…egration

github-actions · 2025-06-16T10:58:26Z

Bijectors.jl documentation for PR #393 is available at:
https://TuringLang.github.io/Bijectors.jl/previews/PR393/

AoifeHughes · 2025-06-16T11:15:17Z

Keeping as draft until my local tests pass

AoifeHughes · 2025-06-16T11:18:47Z

also next commit will remove the Zygote CI

…for EnzymeForward

Copilot

Pull Request Overview

This PR removes Zygote support from Bijectors.jl and updates tests, configuration, and documentation to use ForwardDiff and other supported AD backends.

Deletes the Zygote-specific extension module and removes Zygote entries from configs and tests
Updates test calls and CI workflows to drop Zygote and ensure ForwardDiff/ReverseDiff/Mooncake coverage
Revises documentation examples to use ForwardDiff.gradient instead of Zygote.gradient

Reviewed Changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
test/runtests.jl	Removed `using Zygote`; needs to import `Enzyme` for AD tests
test/ad/utils.jl	Removed Zygote blocks from `test_ad`
test/ad/flows.jl	Added `(:EnzymeForward,)` broken tuple in tests; may be incorrect
test/Project.toml	Deleted Zygote dependency
src/chainrules.jl	Generalized `@debug` comment and added new `rrule` definitions
src/bijectors/pd.jl	Updated comment to drop Zygote reference
ext/BijectorsZygoteExt.jl	Entire file deleted to remove Zygote adjoints
docs/src/examples.md	Swapped out Zygote example for ForwardDiff
docs/Project.toml	Removed Zygote from documentation dependencies
.github/workflows/AD.yml	Removed Zygote from CI matrix

Comments suppressed due to low confidence (3)

src/chainrules.jl:293

New chain rules for aT_b and _vec were added without corresponding tests; consider adding unit tests to validate their gradient behavior.

function ChainRulesCore.rrule(::typeof(aT_b), a::AbstractVector, b::AbstractMatrix)

test/runtests.jl:15

Enzyme is used in AD tests but not imported; add using Enzyme alongside other AD backends

using Tracker

test/ad/flows.jl:3

The second argument to test_ad specifies broken backends, not selection of AD to test. This tuple will not limit to EnzymeForward as intended; remove the (:EnzymeForward,) argument or adjust test_ad to support selecting specific backends.

    test_ad(randn(7), (:EnzymeForward,)) do θ

mhauru

Looks mostly good, one problem with the docs and a couple of questions about PlanarLayer tests as also discussed on Slack.

docs/src/examples.md

mhauru · 2025-06-18T09:17:49Z

src/chainrules.jl

+# Fixes AD issues with `@debug`
 ChainRulesCore.@non_differentiable _debug(::Any)
+
+# ChainRules for utility functions used by PlanarLayer


Are these copied over from somewhere else, or new code? Do we need them? I think ChainRules rules are mostly for Zygote, although other AD packages use them too, so wondering if/why these are necessary.

I think these functions are used by PlanarLayer (src/bijectors/planar_layer.jl). The ChainRules are necessary for proper automatic differentiation support?

ChainRules is only directly used by Zygote, all other (major) AD systems have their own way of handling them. For example ForwardDiff and ReverseDiff use operator overloading. For some AD systems (Mooncake, Enzyme) you can import rules from ChainRules but this is an optional mechanism, because those two backends differentiate code at a lower level. (It's messy, but there's an overview of Julia AD at: https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/faq/differentiability/) So, I'm a bit surprised that we would need to add new rules in this PR, given that we're removing Zygote. Indeed I would have been more expecting to see removal of rules.

If you remove these do any of the tests fail?

(Also, please don't resolve comments unless you're sure that it's been resolved! Otherwise it's hard for people to see what still needs to be discussed)

@AoifeHughes, did you a chance to check yet what happens if you remove these?

Yes, I'm doing it right now, it still breaks :(

Hmm, odd. Happy to have a look at the error and figure it out together if it helps.

Not gonna lie, I'm totally lost and feel like Im just randomly changing things I don't understand here. I'd really appreciate someone taking over so I can learn a bit about how best to work through the problem

Absolutely. Ping me or Penny on Slack when you have a moment and let's have a look at it together.

test/ad/flows.jl

penelopeysm · 2025-06-24T01:14:22Z

I don't think there's a need for me to review this separately. I'm happy to follow up / approve if further changes are made but I think Markus's preexisting comments should be dealt with first, in particular the one about addition of ChainRulesCore functionality.

AoifeHughes · 2025-07-03T09:24:32Z

Im super stuck on this and just cannot get the tests to work. I'd really appreciate any kind of help with this PR.

mhauru · 2025-07-07T11:30:28Z

src/chainrules.jl

+# Fixes AD issues with `@debug`
 ChainRulesCore.@non_differentiable _debug(::Any)


Do you know if this is still needed? Does some other AD backend rely on this?

mhauru · 2025-07-07T11:32:37Z

src/chainrules.jl

+# Fixes AD issues with `@debug`
 ChainRulesCore.@non_differentiable _debug(::Any)
+
+# ChainRules for utility functions used by PlanarLayer


@AoifeHughes, did you a chance to check yet what happens if you remove these?

src/chainrules.jl

mhauru

Thanks @AoifeHughes!

Remove Zygote dependency and update documentation for ForwardDiff int…

160d289

…egration

AoifeHughes requested a review from Copilot June 16, 2025 10:57

AoifeHughes self-assigned this Jun 16, 2025

AoifeHughes added doc dependencies Pull requests that update a dependency file labels Jun 16, 2025

This comment was marked as outdated.

Sign in to view

AoifeHughes added 2 commits June 16, 2025 12:18

Remove Zygote from AD workflow jobs

990da50

fixed failing tests

502952d

AoifeHughes marked this pull request as ready for review June 16, 2025 11:56

AoifeHughes requested a review from Copilot June 16, 2025 12:06

This comment was marked as outdated.

Sign in to view

AoifeHughes added 3 commits June 16, 2025 13:18

bump to remove the expected broken things|

218aab9

tests pls

fd7260c

Add ChainRules for utility functions in PlanarLayer and update tests …

ba1bf9c

…for EnzymeForward

mhauru self-requested a review June 18, 2025 09:05

AoifeHughes requested a review from Copilot June 18, 2025 09:06

Copilot AI reviewed Jun 18, 2025

View reviewed changes

mhauru requested changes Jun 18, 2025

View reviewed changes

yebai requested a review from penelopeysm June 23, 2025 21:33

This comment was marked as resolved.

Sign in to view

penelopeysm removed their request for review July 6, 2025 10:51

AoifeHughes force-pushed the ah/remove-zg-refs branch from 0f5c58d to ba1bf9c Compare July 7, 2025 09:17

mhauru reviewed Jul 7, 2025

View reviewed changes

remove things to try and pass the thing

3ccd657

AoifeHughes force-pushed the ah/remove-zg-refs branch from 8676750 to 3ccd657 Compare July 7, 2025 13:15

github-actions bot reviewed Jul 7, 2025

View reviewed changes

src/chainrules.jl Show resolved Hide resolved

Merge branch 'main' into ah/remove-zg-refs

2d2d684

mhauru approved these changes Jul 7, 2025

View reviewed changes

AoifeHughes merged commit 0f505ed into main Jul 7, 2025
27 of 29 checks passed

AoifeHughes deleted the ah/remove-zg-refs branch July 7, 2025 15:30

		# Fixes AD issues with `@debug`
		ChainRulesCore.@non_differentiable _debug(::Any)

Remove Zygote dependency and update documentation for ForwardDiff integration #393

Remove Zygote dependency and update documentation for ForwardDiff integration #393

Uh oh!

Conversation

AoifeHughes commented Jun 16, 2025

Remove Zygote Support from Bijectors.jl - Addresses: #374

Summary

Changes Made

Files Deleted

Project Configuration

Source Code

Documentation

Impact

Testing

Uh oh!

github-actions bot commented Jun 16, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

AoifeHughes commented Jun 16, 2025

Uh oh!

AoifeHughes commented Jun 16, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jul 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

penelopeysm commented Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

AoifeHughes commented Jul 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Remove Zygote Support from `Bijectors.jl` - Addresses: #374

penelopeysm Jul 1, 2025 •

edited

Loading

penelopeysm commented Jun 24, 2025 •

edited

Loading