model.sanity_check() then jit(ns) causes tracer leak

Joshuaalbert / jaxns

Probabilistic Programming and Nested sampling in JAX

https://jaxns.readthedocs.io/

Other

142 stars 10 forks source link

Closed Joshuaalbert closed 10 months ago

Joshuaalbert commented 10 months ago

Describe the bug Doing a sanity check of model precomputes some things that get interpretted as a context leak when JIT-ing sampler later.

Temporary solution is to comment out model.sanity_check once it passes.

Joshuaalbert commented 10 months ago

Done in #107