Thanks for this comment, I probably could've stressed JIT more. Random question - I noted in the article that you all at DeepMind announced that you're using JA

Thanks for this comment, I probably could've stressed JIT more. Random question - I noted in the article that you all at DeepMind announced that you're using JAX to accelerate your research. IIRC, you guys standardized TensorFlow several years back. What does the current split look like between JAX and TF internally? Do some people use TF and some use JAX, or do you use JAX for specific tasks?