Jax is, as far as an outsider can tell, Google research/DeepMind’s primary computation library now, not Tensorflow. So it’s safer than most Google projects, unless you think they’re secretly developing their third tensor computation library in 10 years (I suppose if anyone was going to do it, it would be Google).