Change to DifferentiationInterface #46

zuhengxu · 2025-02-16T23:16:45Z

Change the optimization interface to DifferentatiationInterface for easier switch/test between different AD systems.

Red-Portal · 2025-02-18T00:44:39Z

Looks good to me!

Project.toml

sunxd3

Thanks for doing this David, so rather minor things.

Project.toml

src/NormalizingFlows.jl

src/optimize.jl

zuhengxu · 2025-03-04T18:42:51Z

As of version 0.4.101 of Mooncake, the problem mentioned above ought to be resolved -- could you please try bumping the version and seeing whether it does?

Thanks @willtebbutt, it works!

ext/NormalizingFlowsEnzymeCoreExt.jl

zuhengxu · 2025-03-05T01:08:23Z

@yebai The NFEnzymeCoreExt is removed. This PR should be ready to be merged. Let me know if I can hit the green button!

src/NormalizingFlows.jl

src/optimize.jl

Co-authored-by: David Widmann <[email protected]>

docs/src/api.md

Co-authored-by: David Widmann <[email protected]>

* switch to differentiationinterface from diffresults * rename train.jl to optimize.jl * fix some compat issue and bump version * update tests to new interface * add Moonkcake to extras * rm all ext for now * rm enzyme test, and import mooncake for test * fixing compat and test with mooncake * fixing test bug * fix _value_and_grad wrapper bug * fix AutoReverseDiff argument typo * minor ed * minor ed * fixing test * minor ed * rm test for mooncake * fix doc * chagne CI * update CI * streamline project toml * Apply suggestions from code review Co-authored-by: David Widmann <[email protected]> * add enzyme to tests * add Enzyme to using list * fixing enzyme readonly error by wrapping loss in Const * mv enzyme related edits to ext/ and fix tests * fixing extension loading error * Update Project.toml Co-authored-by: Guillaume Dalle <[email protected]> * remove Requires Co-authored-by: Guillaume Dalle <[email protected]> * remove explit load ext Co-authored-by: Guillaume Dalle <[email protected]> * Update src/objectives/loglikelihood.jl Co-authored-by: Guillaume Dalle <[email protected]> * make ext dep explicit * rm empty argument specialization for _prep_grad and _value_grad * signal empty rng arg * drop Requires * drop Requires * update test to include mooncake * rm unnecessary EnzymeCoreExt * minor update of readme * typo fix in readme * Update src/NormalizingFlows.jl Co-authored-by: David Widmann <[email protected]> * Update src/NormalizingFlows.jl Co-authored-by: David Widmann <[email protected]> * rm time_elapsed from train_flow * Update docs/src/api.md Co-authored-by: David Widmann <[email protected]> --------- Co-authored-by: Hong Ge <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: Guillaume Dalle <[email protected]>

@require

* added CUDA extension * fixed merge issue * rm extra weakdeps * edit toml * @require cuda ext * minor update tests * try to fix ambiguity * rm tmp test file * wip on randdevice interface * move cuda part in ext * rename sampler.jl to sample.jl * update objectives/elbo.jl to use rand_device * update test/cuda.jl * rm exmaple/test.jl * fix CI test error * fix CI error * update flux compat for test/ * rm tmp test file * fix cuda err diffresults.gradient_result * Documentation (#27) * add NF intro * set up doc files * add gitignore * minor update to readme * update home page * update docs for each funciton * update docs * src * update function docs * update docs * fix readme math rendering issue * update docs * update example doc * update customize layer docs * finish docs * finish docs * Update README.md Co-authored-by: Cameron Pfiffer <[email protected]> * Update README.md Co-authored-by: Cameron Pfiffer <[email protected]> * Update README.md Co-authored-by: Cameron Pfiffer <[email protected]> * Update docs/src/index.md Co-authored-by: Xianda Sun <[email protected]> * Update README.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/index.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/index.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/index.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/index.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/example.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * Update docs/src/customized_layer.md Co-authored-by: Xianda Sun <[email protected]> * minor ed * minor ed to fix latex issue * minor update --------- Co-authored-by: Cameron Pfiffer <[email protected]> Co-authored-by: Xianda Sun <[email protected]> * add more synthetic targets (#20) * add Neal's Funnel and Warped Gaussian * fixed bug in warped gaussian * add reference for warped Gauss * add Cross dsitribution * udpate docs for cross * Update example/targets/cross.jl change comment into docs Co-authored-by: Tor Erlend Fjelde <[email protected]> * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * update cross docs * minor ed * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * doc banana using * fixing docs with latex * baanan docs with latex * add NF quick intro * Revert "add NF quick intro" This reverts commit e399274. * rm unnecesary code for cross * rm example/manifest * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * minor update to cross docs --------- Co-authored-by: Tor Erlend Fjelde <[email protected]> * Fix math display errors in readme (#28) * fix readme math err * minor update readme * add PR links to TODO * CompatHelper: bump compat for Optimisers to 0.3, (keep existing compat) (#29) * CompatHelper: bump compat for Optimisers to 0.3, (keep existing compat) * bump version to 0.1.1 --------- Co-authored-by: CompatHelper Julia <[email protected]> Co-authored-by: Zuheng <[email protected]> * CompatHelper: bump compat for ADTypes to 0.2, (keep existing compat) (#31) Co-authored-by: CompatHelper Julia <[email protected]> * CompatHelper: bump compat for ADTypes to 1, (keep existing compat) (#32) Co-authored-by: CompatHelper Julia <[email protected]> Co-authored-by: Hong Ge <[email protected]> * CompatHelper: bump compat for Enzyme in [weakdeps] to 0.12, (keep existing compat) (#34) Co-authored-by: CompatHelper Julia <[email protected]> * fix doc build failure (#35) * minor doc dep update * rm unnecessary dep * add distributions random linearalgebra to docs env * Create DocNav.yml * CompatHelper: bump compat for Bijectors to 0.14, (keep existing compat) (#41) Co-authored-by: CompatHelper Julia <[email protected]> * Documentation and Turing Navigation CI improvement (#45) * Update CI.yml * Create Docs.yml * Update DocNav.yml * No need of deploydocs() after using new Docs & DocsNav workflows * CompatHelper: bump compat for Enzyme in [weakdeps] to 0.13, (keep existing compat) (#39) Co-authored-by: CompatHelper Julia <[email protected]> * Change to DifferentiationInterface (#46) * switch to differentiationinterface from diffresults * rename train.jl to optimize.jl * fix some compat issue and bump version * update tests to new interface * add Moonkcake to extras * rm all ext for now * rm enzyme test, and import mooncake for test * fixing compat and test with mooncake * fixing test bug * fix _value_and_grad wrapper bug * fix AutoReverseDiff argument typo * minor ed * minor ed * fixing test * minor ed * rm test for mooncake * fix doc * chagne CI * update CI * streamline project toml * Apply suggestions from code review Co-authored-by: David Widmann <[email protected]> * add enzyme to tests * add Enzyme to using list * fixing enzyme readonly error by wrapping loss in Const * mv enzyme related edits to ext/ and fix tests * fixing extension loading error * Update Project.toml Co-authored-by: Guillaume Dalle <[email protected]> * remove Requires Co-authored-by: Guillaume Dalle <[email protected]> * remove explit load ext Co-authored-by: Guillaume Dalle <[email protected]> * Update src/objectives/loglikelihood.jl Co-authored-by: Guillaume Dalle <[email protected]> * make ext dep explicit * rm empty argument specialization for _prep_grad and _value_grad * signal empty rng arg * drop Requires * drop Requires * update test to include mooncake * rm unnecessary EnzymeCoreExt * minor update of readme * typo fix in readme * Update src/NormalizingFlows.jl Co-authored-by: David Widmann <[email protected]> * Update src/NormalizingFlows.jl Co-authored-by: David Widmann <[email protected]> * rm time_elapsed from train_flow * Update docs/src/api.md Co-authored-by: David Widmann <[email protected]> --------- Co-authored-by: Hong Ge <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: Guillaume Dalle <[email protected]> * add CUDA extension * undo change * add CUDA to test dep * add `rand_device` in main package * fix test error * move CUDA test to separate folder for gpu testing * add some preliminary gpu test CI code * remove cuda.jl from testing * update gpu test pipeline setup * update test file path * try fix error in CI setup * move files * try to debug * fix error * fix cuda testing project.toml * cosmetic fixed -- resolve linter complaints * fix test error * clean up * fix more errors * try to fix test error * add back allowscalar * refactoring * try fixing doc error * try fix doc again * add ref link --------- Co-authored-by: Zuheng <[email protected]> Co-authored-by: David Xu <[email protected]> Co-authored-by: Cameron Pfiffer <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: CompatHelper Julia <[email protected]> Co-authored-by: Hong Ge <[email protected]> Co-authored-by: Shravan Goswami <[email protected]> Co-authored-by: Xianda Sun <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: Guillaume Dalle <[email protected]>

zuhengxu added 18 commits February 16, 2025 15:12

switch to differentiationinterface from diffresults

6220e96

rename train.jl to optimize.jl

7b4fb85

fix some compat issue and bump version

5ff6041

update tests to new interface

3010669

add Moonkcake to extras

f7ee84b

rm all ext for now

e68ef5f

rm enzyme test, and import mooncake for test

9bdf1f7

fixing compat and test with mooncake

b7f9f08

fixing test bug

1970b09

fix _value_and_grad wrapper bug

b0390f7

fix AutoReverseDiff argument typo

3b40c33

minor ed

9552de1

minor ed

104e5dd

fixing test

d1ec834

minor ed

deba738

rm test for mooncake

8976307

fix doc

cb3db4a

chagne CI

1615009

zuhengxu requested review from sunxd3, yebai, Red-Portal and devmotion February 17, 2025 04:32

yebai requested a review from mhauru February 20, 2025 11:21

This was referenced Feb 20, 2025

Fix Enzyme extension #38

Closed

CompatHelper: bump compat for Zygote in [weakdeps] to 0.7, (keep existing compat) #44

Closed

yebai approved these changes Feb 20, 2025

View reviewed changes

yebai reviewed Feb 20, 2025

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

Merge branch 'main' into diffinterface

690f754

sunxd3 requested changes Feb 20, 2025

View reviewed changes

zuhengxu added 6 commits March 4, 2025 10:39

make ext dep explicit

3e65cde

rm empty argument specialization for _prep_grad and _value_grad

eeb9a92

signal empty rng arg

4b97585

drop Requires

9075266

drop Requires

eade7a3

update test to include mooncake

91202ff

yebai reviewed Mar 4, 2025

View reviewed changes

ext/NormalizingFlowsEnzymeCoreExt.jl Outdated Show resolved Hide resolved

yebai reviewed Mar 4, 2025

View reviewed changes

ext/NormalizingFlowsEnzymeCoreExt.jl Outdated Show resolved Hide resolved

zuhengxu requested review from yebai and gdalle March 4, 2025 21:06

rm unnecessary EnzymeCoreExt

dcee3c0

zuhengxu mentioned this pull request Mar 5, 2025

CompatHelper: bump compat for Zygote in [weakdeps] to 0.7, (keep existing compat) #47

Closed

zuhengxu added 2 commits March 4, 2025 17:18

minor update of readme

edf2f12

typo fix in readme

9fc328d

yebai approved these changes Mar 5, 2025

View reviewed changes

devmotion reviewed Mar 5, 2025

View reviewed changes

src/NormalizingFlows.jl Outdated Show resolved Hide resolved

src/NormalizingFlows.jl Outdated Show resolved Hide resolved

src/optimize.jl Show resolved Hide resolved

zuhengxu and others added 3 commits March 5, 2025 08:39

Update src/NormalizingFlows.jl

41d10ef

Co-authored-by: David Widmann <[email protected]>

Update src/NormalizingFlows.jl

db4872b

Co-authored-by: David Widmann <[email protected]>

rm time_elapsed from train_flow

0fe536f

zuhengxu requested a review from devmotion March 5, 2025 17:51

devmotion reviewed Mar 5, 2025

View reviewed changes

docs/src/api.md Outdated Show resolved Hide resolved

Update docs/src/api.md

fc935ce

Co-authored-by: David Widmann <[email protected]>

yebai merged commit 4a018d8 into main Mar 5, 2025
3 checks passed

yebai deleted the diffinterface branch March 5, 2025 20:57

yebai mentioned this pull request Mar 5, 2025

add compiled tape support for ReverseDiff.jl #23

Closed

yebai mentioned this pull request Mar 13, 2025

Replace ForwardDiff/ReverseDiff/Zygote/Mooncake/Enzyme gradient call with DifferentiationInterface. TuringLang/Bijectors.jl#373

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Change to DifferentiationInterface #46

Change to DifferentiationInterface #46

Uh oh!

zuhengxu commented Feb 16, 2025

Uh oh!

Red-Portal commented Feb 18, 2025

Uh oh!

Uh oh!

sunxd3 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zuhengxu commented Mar 4, 2025

Uh oh!

Uh oh!

Uh oh!

zuhengxu commented Mar 5, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Change to DifferentiationInterface #46

Change to DifferentiationInterface #46

Uh oh!

Conversation

zuhengxu commented Feb 16, 2025

Uh oh!

Red-Portal commented Feb 18, 2025

Uh oh!

Uh oh!

sunxd3 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zuhengxu commented Mar 4, 2025

Uh oh!

Uh oh!

Uh oh!

zuhengxu commented Mar 5, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!