JAX still has the

"This is a research project, not an official Google product. Expect bugs and sharp edges. Please help by trying it out, reporting bugs, and letting us know what you think!"

disclaimer in its readme. This is quite scary, especially coming from Google which is known to abandon projects out of the blue