probcomp / Gen.jl

A general-purpose probabilistic programming system with programmable inference
https://gen.dev
Apache License 2.0
1.79k stars 160 forks source link

Use Zygote.jl #128

Open marcoct opened 4 years ago

marcoct commented 4 years ago

Starting looking into this here: https://github.com/probcomp/Gen/blob/20190927-marcoct-zygote/zygote/zygote_test.jl

Ran into https://github.com/FluxML/Zygote.jl/issues/347

ianfiske commented 4 years ago

Just a heads up that a fix for FluxML/Zygote.jl#347 has been merged into Zygote.

ludgerpaehler commented 4 years ago

Hi, has there been any progress on this sub-project? If not, I'd love to volunteer.

marcoct commented 4 years ago

Over the past few months I've been working on this intermittently on this branch: https://github.com/probcomp/Gen/tree/20190927-marcoct-zygote. Today, I finished a prototype Zygote port for AD in the dynamic modeling language. It passes the existing unit tests, but is not at all optimized, and anecdotally based on speed of unit tests it looks like there are some performance issues.

Remaining work is to (i) benchmark and stress testing to characterize any performance surprises relative to the older version, (ii) profile and optimize or improve the implementation, and (iii) port it for the static modeling language as well.

@ludgerpaehler Thank you for volunteering! If you are interested in doing any benchmarks of code that uses choice_gradients or accumulate_param_gradients!, comparing this branch versus master, that would be helpful. In particular I suspect that the test in Gen/test/inference/variational.jl became much slower.

ludgerpaehler commented 4 years ago

@marcoct I'll get started on (i) and (ii).