Fix caching bug.
Fix caching bug.

This was a very tricky bug. The compilation caching keys depend on
many pieces of the transformation state, including among other
things, the pointer-value of the Jaxpr. Since in the previous
commit I have added a `convert_constvars_jaxpr` call, every time
that one ran, it produces the same semantic Jaxpr but a different
pointer value, which was breaking caching.

The fix is to cache the `convert_constvars_jaxpr`.

google/jax,google/jax,tensorflow/probability,google/jax,tensorflow/probability,google/jax