def setUp(self):
    super().setUp()
    jax._src.lax.control_flow._initial_style_open_jaxpr.cache_clear()
    jax._src.lax.control_flow._initial_style_jaxpr.cache_clear()
    jax._src.lax.control_flow._initial_style_jaxprs_with_common_consts.cache_clear()
def setUp(self):
    super().setUp()
    jax._src.lax.control_flow._initial_style_open_jaxpr.cache_clear()
    jax._src.lax.control_flow._initial_style_jaxpr.cache_clear()
    jax._src.lax.control_flow._initial_style_jaxprs_with_common_consts.cache_clear()