GoogleJAX是一個(gè)用于變換數(shù)值函數(shù)的機(jī)器學(xué)習(xí)框架，Google稱其為為結(jié)合了修改版本的Autograd（通過函數(shù)微分自動(dòng)獲得梯度函數(shù)）和TensorFlow的XLA（加速線性代數(shù)）。該框架的設(shè)計(jì)盡可能遵循NumPy的結(jié)構(gòu)和工作流程，并與TensorFlow和PyTorch等各種現(xiàn)有框架協(xié)同工作。

JAX的主要功能是包括：

• grad：自動(dòng)微分
• jit：編譯
• vmap：自動(dòng)矢量化
• pmap：SPMD編程