JAX seems great but the Google ghost is a lot to stomach. The risk of JAX getting axed or replaced with a JAX 2.0 - completely incompatible with existing JAX code - is not insignificant.