Speed is one reason why Google is using Jax. With decorators, vmap/pmap it jit compiles your Python/Numpy code for fast/vectorized execution on GPU, TPU and CPU
And Instagram have their own CPython fork called Cinder - https://github.com/facebookincubator/cinder