Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Basic rewrite of the package 2023 edition Part I: ADVI (#49)
* refactor ADVI, change gradient operation interface * remove unused file, remove unused dependency * fix ADVI elbo computation more efficiently * fix missing entropy regularization term * add LogDensityProblem interface * refactor use bijectors directly instead of transformed distributions This is to avoid having to reconstruct transformed distributions all the time. The direct use of bijectors also avoids going through lots of abstraction layers that could break. Instead, transformed distributions could be constructed only once when returing the VI result. * fix type restrictions * remove unused file * fix use of with_logabsdet_jacobian * restructure project; move the main VI routine to its own file * remove redundant import * restructure project into more modular objective estimators * migrate to AbstractDifferentiation * add location scale pre-packaged variational family, add functors * Revert "migrate to AbstractDifferentiation" This reverts commit 2a4514e. * fix use optimized MvNormal specialization, add logpdf for Loc.Scale. * remove dead code * fix location-scale logpdf - Full Monte Carlo ELBO estimation now works. I checked. * add sticking-the-landing (STL) estimator * migrate to Optimisers.jl * remove execution time measurement (replace later with somethin else) * fix use multiple dispatch for deciding whether to stop entropy grad. * add termination decision, callback arguments * add Base.show to modules * add interface calling `restructure`, rename rebuild -> restructure * add estimator state interface, add control variate interface to ADVI * fix `show(advi)` to show control variate * fix simplify `show(advi.control_variate)` * fix type piracy by wrapping location-scale bijected distribution * remove old AdvancedVI custom optimizers * fix Location Scale to not depend on Bijectors * fix RNG namespace * fix location scale logpdf bug * add Accessors dependency * add location scale, autodiff tests * add Accessors import statement * remove optimiser tests * refactor slightly generalize the distribution tests for the future * migrate to SimpleUnPack, migrate to ADTypes * rename vi.jl to optimize.jl * fix estimate_gradient to use adtypes * add exact inference tests * remove Turing dependency in tests * remove unused projection * remove redundant `ADVIEnergy` object (now baked into `ADVI`) * add more tests, fix rng seed for tests * add more tests, fix seed for tests * fix non-determinism bug * fix test hyperparameters so that tests pass, minor cleanups * fix minor reorganization * add missing files * fix add missing file, rename adbackend argument * fix errors * rename test suite * refactor renamed arguments for ADVI to be shorter * fix compile error in advi test * add initial doc * remove unused epsilon argument in location scale * add project file for documenter * refactor STL gradient calculation to use multiple dispatch * fix type bugs, relax test threshold for the exact inference tests * refactor derivative utils to match NormalizingFlows.jl with extras * add documentation, refactor optimize * fix bug missing extension * remove tracker from tests * remove export for internal derivative utils * fix test errors, old interface * fix wrong derivative interface, add documentation * update documentation * add doc build CI * remove convergence criterion for now * remove outdated export * update documentation * update documentation * update documentation * fix type error in test * remove default ADType argument * update README * update make getting started example actually run Julia * fix remove Float32 tests for inference tests * update version * add documentation publishing url * fix wrong uuid for ForwardDiff * Update CI.yml * refactor use `sum` and `mean` instead of abusing `mapreduce` * remove tests for `FullMonteCarlo` * add tests for the `optimize` interface * fix turn off Zygote tests for now * remove unused function * refactor change bijector field name, simplify STL estimator * update documentation * update STL documentation * update STL documentation * update location scale documentation * fix README * fix math in README * add gradient to arguments of callback!, remove `gradient_norm` info * fix math in README.md Co-authored-by: David Widmann <[email protected]> * fix type constraint in `ZygoteExt` Co-authored-by: David Widmann <[email protected]> * fix import of `Random` Co-authored-by: David Widmann <[email protected]> * refactor `__init__()` Co-authored-by: David Widmann <[email protected]> * fix type constraint in definition of `value_and_gradient!` Co-authored-by: David Widmann <[email protected]> * refactor `ZygoteExt`; use `only` instead of `first` Co-authored-by: David Widmann <[email protected]> * refactor type constraint in `ReverseDiffExt` Co-authored-by: David Widmann <[email protected]> * refactor remove outdated debug mode macro * fix remove outdated DEBUG mechanism * fix LaTeX in README: `operatorname` is currently broken * remove `SimpleUnPack` dependency * fix LaTeX in docs and README * add warning about forward-mode AD when using `LocationScale` * fix documentation * fix remove reamining use of `@unpack` * Revert "remove `SimpleUnPack` dependency" This reverts commit 29d7d27. * Revert "fix remove reamining use of `@unpack`" This reverts commit 8173744. * fix documentation for `optimize` * add specializations of `Optimise.destructure` for mean-field * This fixes the poor performance of `ForwardDiff` * This prevents the zero elements of the mean-field scale being extracted * add test for `Optimisers.destructure` specializations * add specialization of `rand` for meanfield resulting in faster AD * add argument checks for `VIMeanFieldGaussian`, `VIFullRankGaussian` * update documentation * fix type instability, bug in argument check in `LocationScale` * add missing import bug * refactor test, fix type bug in tests for `LocationScale` * add missing compat entries * fix missing package import in test * add additional tests for sampling `LocationScale` * fix bug in batch in-place `rand!` for `LocationScale` * fix bug in inference test initialization * add missing file * fix remove use of for 1.6 * refactor adjust inference test hyperparameters to be more robust * refactor `optimize` to return `obj_state`, add warm start kwargs * refactor make tests more robust, reduce amount of tests * fix remove a cholesky in test model * fix compat bounds, remove unused package * bump compat for ADTypes 0.2 * fix broken LaTeX in README * remove redundant use of PDMats in docs * fix use `Cholesky` signature supported in 1.6 * revert custom variational families and docs * remove doc action for now * revert README for now * refactor remove redundant `rng` argument to `ADVI`, improve docs * fix wrong whitespace in tests * refactor `estimate_gradient` to `estimate_gradient!`, add docs * refactor add default `init` impl, update docs * merge (manually) commit ff32ac6 * fix test for new interface, change interface for `optimize`, `advi` * fix integer subtype error in documentation of advi Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix remove redundant argument for `advi` * remove manifest * refactor remove imports and use fully qualified names * update documentation for `AbstractVariationalObjective` Co-authored-by: Tor Erlend Fjelde <[email protected]> * refactor use StableRNG instead of Random123 * refactor migrate to Test, re-enable x86 tests * refactor remove inner constructor for `ADVI` * fix swap `export`s and `include`s Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix doscs for `ADVI` Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix use `FillArrays` in the test problems Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix `optimize` docs Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix improve argument names and docs for `optimize` * fix tests to match new interface of `optimize` * refactor move utility functions to new file * fix docs for `optimize` Co-authored-by: Tor Erlend Fjelde <[email protected]> * refactor advi internal objective Co-authored-by: Tor Erlend Fjelde <[email protected]> * refactor move `rng` to be an optional first argument * fix docs for optimize * add compat bounds to test dependencies * update compat bound for `Optimisers` * fix test compat * fix remove `!` in callback Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix rng argument position in `advi` * fix callback signature in `optimize` * refactor reorganize test files and naming * fix simplify description for `optimize` Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix remove redundant `Nothing` type signature for `maybe_init` * fix remove "internal use" warning in documentation * refactor change `estimate_gradient!` signature to be type stable * add signature for computing `advi` over a fixed set of samples * fix change test tolerance * fix update documentation for `estimate_gradient!` * refactor remove type constraint for variational parameters * fix remove dead code * add compat entry for stdlib * add compat entry for stdlib in `test/` * fix rng argument position in tests * refactor change name of inference test * fix documentation for `optimize` * refactor rewrite the documentation for the global interfaces * fix compat error * fix documentation for `optimize` to be single line * refactor remove begin end for one-liner * refactor create unified interface for estimating objectives * refactor unify interface for entropy estimator, fix advi docs * fix STL estimator to use manually stopped gradients instead * add inference test for a non-bijector model * refactor add indirections to handle STL and bijectors in ADVI * refactor split inference tests for advi+distributionsad * refactor rename advi to repgradelbo and not use bijectors directly * fix documentation for estimate_objective * refactor add indirection in repgradelbo for interacting with `q` * add TransformedDistribution support as extension * Update src/objectives/elbo/repgradelbo.jl Co-authored-by: Tor Erlend Fjelde <[email protected]> * fix docstring for entropy estimator * fix `reparam_with_entropy` specialization for bijectors * enable Zygote for non-bijector tests --------- Co-authored-by: Hong Ge <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: Tor Erlend Fjelde <[email protected]>
- Loading branch information
4 people
authored
Dec 8, 2023
1 parent
b0c4be3
commit 576259a