Both DeepMind and Brain use Jax, so they will definitely use Jax. However, they use different high level frameworks: All of DeepMind uses Haiku, while on the Brain side there are competing frameworks, with flax currently being the most often used one AFAIK. I'm not aware of anyone using trax there, and I would not expect it to get more adoption, on the contrary.