imported>Mhss：/* top */

2026-05-30T17:23:35Z

top

新页面

{{Infobox software
| title = JAX
| name =
| logo = Google JAX logo.svg
| logo caption =
| logo alt =
| logo size = 200px
| collapsible = 
| screenshot =
| screenshot size =
| screenshot alt =
| caption =
| author =
| developer = [[Google]], [[Nvidia]]<ref>{{cite web |url=https://github.com/jax-ml/jax/blob/main/AUTHORS |title=jax/AUTHORS at main · jax-ml/jax |date= |author= |website=[[GitHub]] |accessdate= December 21, 2024}}</ref>
| released = {{Start date and age|2019|10|31|df=yes}}<ref>{{cite web|url=https://github.com/google/jax/releases/tag/jax-v0.1.49|title=jax-v0.1.49}}</ref>
| latest release version =
| latest release date = 
| repo = {{URL|https://github.com/jax-ml/jax}}
| programming language = [[Python]], [[C++]]
| middleware =
| operating system = [[Linux]], [[macOS]], [[Windows]]
| platform = [[Python]], [[NumPy]]
| size =
| language count = 
| language footnote =
| genre = [[机器学习]]
| license = [[Apache 2.0]]
| website = {{URL|https://docs.jax.dev/en/latest/}}
}}
'''JAX'''，由[[Google]]开发、並由[[Nvidia]]做出部分貢獻<ref name=":0">
{{Citation |title=JAX: Autograd and XLA |date=2022-06-18 |url=https://github.com/google/jax |archive-url=https://web.archive.org/web/20220618205214/https://github.com/google/jax |publisher=Google |bibcode=2021ascl.soft11002B |access-date=2022-06-18 |archive-date=2022-06-18|last1=Bradbury |first1=James |last2=Frostig |first2=Roy |last3=Hawkins |first3=Peter |last4=Johnson |first4=Matthew James |last5=Leary |first5=Chris |last6=MacLaurin |first6=Dougal |last7=Necula |first7=George |last8=Paszke |first8=Adam |last9=Vanderplas |first9=Jake |last10=Wanderman-Milne |first10=Skye |last11=Zhang |first11=Qiao |journal=Astrophysics Source Code Library }}</ref><ref>{{Cite journal |last1=Frostig |first1=Roy |last2=Johnson |first2=Matthew James |last3=Leary |first3=Chris |date=2018-02-02 |year=2018 |title=Compiling machine learning programs via high-level tracing |url=https://mlsys.org/Conferences/doc/2018/146.pdf |url-status=live |journal=MLsys |pages=1–3 |archive-url=https://web.archive.org/web/20220621153349/https://mlsys.org/Conferences/doc/2018/146.pdf |archive-date=2022-06-21}}</ref><ref>
{{Cite web |title=Using JAX to accelerate our research |url=https://www.deepmind.com/blog/using-jax-to-accelerate-our-research |url-status=live |archive-url=https://web.archive.org/web/20220618205746/https://www.deepmind.com/blog/using-jax-to-accelerate-our-research |archive-date=2022-06-18 |access-date=2022-06-18 |website=www.deepmind.com |language=en}}</ref>的[[Python]][[机器学习]]框架，用於变换数值函数。JAX结合了修改版的Autograd自動微分系統<ref>{{Cite web |url=https://github.com/HIPS/autograd |title=autograd |access-date=2023-09-23 |archive-date=2022-07-18 |archive-url=https://web.archive.org/web/20220718131101/https://github.com/hips/autograd |dead-url=no }}</ref>，以及來自OpenXLA專案的編譯器{{en-link|加速线性代数|Accelerated Linear Algebra|XLA}}<ref>{{Cite web |url=https://www.tensorflow.org/xla |title=XLA |access-date=2023-09-23 |archive-date=2022-09-01 |archive-url=https://web.archive.org/web/20220901124303/https://www.tensorflow.org/xla |dead-url=no }}</ref>，可加速數值線性運算。其設計目標是在介面與程式設計風格上盡可能與[[NumPy]]保持相容，使使用者能夠以熟悉的方式撰寫高效能運算程式。此外，JAX亦可與[[TensorFlow]]、[[PyTorch]]等機器學習框架整合使用。<ref>{{Cite web |last=Lynley |first=Matthew |title=Google is quietly replacing the backbone of its AI product strategy after its last big push for dominance got overshadowed by Meta |url=https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6 |archive-url=https://web.archive.org/web/20220621143905/https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6 |archive-date=2022-06-21 |access-date=2022-06-21 |website=Business Insider |language=en-US}}</ref><ref>{{Cite web |date=2022-04-25 |title=Why is Google's JAX so popular? |url=https://analyticsindiamag.com/why-is-googles-jax-so-popular/ |url-status=live |archive-url=https://web.archive.org/web/20220618210503/https://analyticsindiamag.com/why-is-googles-jax-so-popular/ |archive-date=2022-06-18 |access-date=2022-06-18 |website=Analytics India Magazine |language=en-US}}</ref>

==主要功能==
JAX的主要功能是<ref name=":0" />：
*grad：自动微分，
*jit：即时编译，
*vmap：自动向量化，
*pmap：{{en-link|单程序多数据|Single program, multiple data|SPMD}}编程。

== grad ==
{{Main|自动微分}}
下面的代码演示{{code|grad}}函数的自动微分。

<syntaxhighlight lang="numpy">
# 导入库
from jax import grad
import jax.numpy as jnp

# 定义logistic函数
def logistic(x):
return jnp.exp(x) / (jnp.exp(x) + 1)

# 获得logistic函数的梯度函数
grad_logistic = grad(logistic)

# 求值logistic函数在x = 1处的梯度
grad_log_out = grad_logistic(1.0)
print(grad_log_out)
</syntaxhighlight>

最终的输出为：
<syntaxhighlight lang="output">
0.19661194
</syntaxhighlight>

== jit ==
{{Main|即时编译}}
下面的代码演示<code>jit</code>函数的优化。

<syntaxhighlight lang="numpy">
# 导入库
from jax import jit
import jax.numpy as jnp

# 定义cube函数
def cube(x):
return x * x * x

# 生成数据
x = jnp.ones((10000, 10000))

# 创建cube函数的jit版本
jit_cube = jit(cube)

# 应用cube函数和jit_cube函数于相同数据来比较其速度
cube(x)
jit_cube(x)
</syntaxhighlight>

可见{{code|jit_cube}}的运行时间显著的短于{{code|cube}}。

== vmap ==
{{Main|{{en-link|自动向量化|Automatic vectorization}}}}
下面的代码展示{{code|vmap}}函数的通过[[SIMD]]的向量化。

<syntaxhighlight lang="numpy">
# 导入库
from functools import partial
from jax import vmap
import jax.numpy as jnp

# 定义函数
def grads(self, inputs):
in_grad_partial = partial(self._net_grads, self._net_params)
grad_vmap = vmap(in_grad_partial)
rich_grads = grad_vmap(inputs)
flat_grads = np.asarray(self._flatten_batch(rich_grads))
assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
return flat_grads
</syntaxhighlight>

== 使用JAX的库 ==
一些Python库使用JAX作为后端，这包括：
{{div col|2}}
* Flax，最初由[[Google Brain]]开发的高层[[人工神经网络]]库<ref>{{Citation |title=Flax: A neural network library and ecosystem for JAX designed for flexibility |date=2022-07-29 |url=https://github.com/google/flax |publisher=Google |access-date=2022-07-29 |archive-date=2022-09-03 |archive-url=https://web.archive.org/web/20220903103053/https://github.com/google/flax |dead-url=no }}</ref>。
* Equinox，将参数化函数（包括[[人工神经网络]]）表示为PyTree的库。它由Patrick Kidger创建<ref>{{Citation |last=Kidger |first=Patrick |title=Equinox |date=2022-07-29 |url=https://github.com/patrick-kidger/equinox |access-date=2022-07-29 |archive-date=2023-09-19 |archive-url=https://web.archive.org/web/20230919020900/https://github.com/patrick-kidger/equinox |dead-url=no }}</ref>。
* Diffrax，用于求[[常微分方程的数值方法|微分方程的数值解]]的库，比如解[[常微分方程]]和[[随机微分方程]]<ref>{{Citation |last=Kidger |first=Patrick |title=Diffrax |date=2023-08-05 |url=https://github.com/patrick-kidger/diffrax |access-date=2023-08-08 |archive-date=2023-08-10 |archive-url=https://web.archive.org/web/20230810231524/https://github.com/patrick-kidger/diffrax |dead-url=no }}</ref>。
* Optax，[[DeepMind]]开发的用于梯度处理和[[最优化]]的库<ref>{{Citation |title=Optax |date=2022-07-28 |url=https://github.com/deepmind/optax |publisher=DeepMind |access-date=2022-07-29 |archive-date=2023-06-07 |archive-url=https://web.archive.org/web/20230607072614/https://github.com/deepmind/optax |dead-url=no }}</ref>。
* Lineax，用于解[[线性方程组]]和{{en-link|线性最小二乘法的数值方法|Numerical methods for linear least squares|线性最小二乘法}}<ref>{{Citation |title=Lineax |date=2023-08-08 |url=https://github.com/google/lineax |access-date=2023-08-08 |publisher=Google |archive-date=2023-08-10 |archive-url=https://web.archive.org/web/20230810231532/https://github.com/google/lineax |dead-url=no }}</ref>。
* RLax，[[DeepMind]]开发的用于[[强化学习]]的库<ref>{{Citation |title=RLax |date=2022-07-29 |url=https://github.com/deepmind/rlax |publisher=DeepMind |access-date=2022-07-29 |archive-date=2023-04-26 |archive-url=https://web.archive.org/web/20230426101001/https://github.com/deepmind/rlax |dead-url=no }}</ref>
* jraph，DeepMind开发的{{en-link|图神经网络|Graph neural network}}库<ref>{{Citation |title=Jraph - A library for graph neural networks in jax. |date=2023-08-08 |url=https://github.com/deepmind/jraph |access-date=2023-08-08 |publisher=DeepMind |archive-date=2022-11-23 |archive-url=https://web.archive.org/web/20221123231431/https://github.com/deepmind/jraph/ |dead-url=no }}</ref>。
* jaxtyping，用于为阵列或张量的形状和数据类型增加[[类型签名|类型标注]]的库<ref>{{Citation |title=jaxtyping |date=2023-08-08 |url=https://github.com/google/jaxtyping |access-date=2023-08-08 |publisher=Google |archive-date=2023-08-10 |archive-url=https://web.archive.org/web/20230810231646/https://github.com/google/jaxtyping |dead-url=no }}</ref>。
*NumPyro，[[概率编程]]库<ref>{{cite web|url=https://github.com/pyro-ppl/numpyro|title=NumPyro － Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU|access-date=2022-08-31|archive-date=2022-08-31|archive-url=https://web.archive.org/web/20220831135823/https://github.com/pyro-ppl/numpyro|dead-url=no}}</ref>。
*Brax，[[物理引擎]]<ref>{{cite web|url=https://github.com/google/brax|title=Brax － Massively parallel rigidbody physics simulation on accelerator hardware|access-date=2022-08-31|archive-date=2022-08-31|archive-url=https://web.archive.org/web/20220831013301/https://github.com/google/brax|dead-url=no}}</ref>。
{{div col end}}

== 参见 ==
* [[NumPy]]
* [[TensorFlow]]
* [[PyTorch]]
* [[CUDA]]
* [[自动微分]]
* [[即时编译]]
* {{en-link|自动向量化|Automatic vectorization}}
* {{en-link|自动并行|Automatic parallelization}}

== 引用 ==
{{reflist|2}}

== 外部链接 ==
* Documentationː {{URL|https://jax.readthedocs.io/}}
* Colab ([[Jupyter]]/iPython) Quickstart Guideː {{URL|https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/quickstart.ipynb
}}
* [[TensorFlow]]'s XLAː {{URL|https://www.tensorflow.org/xla}} (Accelerated Linear Algebra)
* {{YouTube|id=WdTeDXsOSj4|title=Intro to JAX: Accelerating Machine Learning research}}
* Original paperː {{URL|https://mlsys.org/Conferences/doc/2018/146.pdf}}

{{Differentiable computing}}

[[Category:数据挖掘和机器学习软件]]
[[Category:机器学习]]
[[Category:Google軟體]]
[[Category:用Python編程的自由軟體]]
[[Category:Python科学库]]
[[Category:使用Apache许可证的软件]]
[[Category:开源人工智能]]

JAX - 版本历史

imported>Mhss：​/* top */

imported>Mhss：/* top */