I don't think JAX is a new language though, if it is, it's still Python?

It’s not even really a language as much as it is an interop layer of sorts. It’s a collection of functional operations on numpy arrays, and a capable JIT. It’s main draw is honestly autograd, and if you’re using that, you often can’t use while or for loops effectively, can’t index numpy arrays without .at[], and should avoid side effects, so strangely it’s neither it’s own language or Python any longer