[compatible with jax 0.4.36] remove ensure_compile_time_eval context, fix progress bar in loop transformation

[compatible with jax 0.4.36] remove ensure_compile_time_eval context, fix progress bar in loop transformation #249