flax

Access thread_resources via jax.interpreters.pxla instead of jax.experimental.maps

#3775

Merged