JAX公式ドキュメントのFAQによると、CPU上で個々の配列演算を対象としたマイクロベンチマークでは、演算ごとのディスパッチのオーバーヘッドが小さいNumPyの方が一般的に高速であるとされています。一方で、GPUやTPU上でコードを実行する場合や、CPU上でもより複雑なJITコンパイル済みの一連の演算をベンチマークする場合には、一般的にJAXの方がNumPyを上回る性能を発揮すると説明されています。NumPyが常にCPU上で同期的・逐次的に演算を実行するのに対し、JAXはCPU・GPU・TPUのいずれでも実行可能で、非同期的にディスパッチされる設計を持つため、両者を単純に比較することは難しいとされ、公式ドキュメントは実際の用途に応じて自分自身でベンチマークを行うことを推奨しています。