Always restore _jit_is_disabled
Always restore _jit_is_disabled

Otherwise if you get an exception inside a ``disable_jit()`` context (not
uncommon if debugging), ``jit`` is disabled permanently!

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