<!--lint ignore double-link-->

Awesome JAX <img src="https://yellow-cdn.veclightyear.com/35dd4d3f/187f3dbd-536e-408c-ba05-8557b9ca883e.png" alt="JAX Logo" align="right" height="100">

<!--lint ignore double-link-->

JAX 通过使用类似 NumPy 的 API 将自动微分和 XLA 编译器 集成在一起，以实现 GPU 和 TPU 等加速器上高性能的机器学习研究。

<!--lint enable double-link-->

这是一个精心整理的 JAX 库、项目及其他资源的优秀列表。欢迎任何贡献！

目录

• 库
• 模型和项目
• 视频
• 论文
• 教程和博客文章
• 书籍
• 社区

<a name="libraries" />

库

• 神经网络库
  • Flax - 注重灵活性和清晰性。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/3269125c-5d12-44b5-b530-81560c5e6a85.png" align="center">
  • Haiku - 由 DeepMind 的 Sonnet 作者创建，注重简单性。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/15318852-20bd-4f49-9316-762cc3c3a4c4.png" align="center">
  • Objax - 拥有类似于 PyTorch 的面向对象设计。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/fa2fac6b-7ef5-465e-ad20-4a8f96062871.png" align="center">
  • Elegy - JAX 中的高级深度学习 API，支持 Flax、Haiku 和 Optax。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/c9b5abe5-950f-4483-919e-d7997b8b8ccb.png" align="center">
  • Trax - 提供常见工作负载解决方案的“内置电池”深度学习库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/830ac7b1-8d1a-4df4-a9b3-421641e018a0.png" align="center">
  • Jraph - 轻量级图神经网络库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/3b7b5251-0a20-48f5-a83c-e9e558c62310.png" align="center">
  • Neural Tangents - 用于指定有限和_无限_宽度神经网络的高级 API。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/e68e6920-0b0b-4697-8f79-39aa0a483a2e.png" align="center">
  • HuggingFace - 为广泛的自然语言任务提供预训练的 Transformers 生态系统（Flax）。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/db417fe5-cba3-4e4a-9563-5848681ba4e1.png" align="center">
  • Equinox - 可调用的 PyTrees 和过滤 JIT/grad 转换 => JAX 中的神经网络。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/b942680a-13be-4b74-bf55-ab2af0eb32ff.png" align="center">
  • Scenic - 一个用于计算机视觉研究及其扩展的 JAX 库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/949a999f-de25-4593-af17-1bc0694c7288.png" align="center">
• Levanter - 具备名称张量和 JAX 的易懂、可扩展、可复制的基础模型。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/a642a688-eb43-4834-aa61-9c689f89a2ac.png" align="center">
• EasyLM - 简化 LLMs 的预训练、微调、评估和服务（基于 JAX/Flax 实现）。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/895d34db-225c-437e-8b2b-15e7b5a92958.png" align="center">
• NumPyro - 基于 Pyro 库的概率编程。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/05fbcbf4-fb81-4c65-9cab-96f70009844f.png" align="center">
• Chex - 编写和测试高可靠性 JAX 代码的工具库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/c4e9cd2d-bd8a-4f9b-8dc3-0028083dd32a.png" align="center">
• Optax - 梯度处理和优化库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/1cd6baa5-814f-464a-86e6-d795d57eb45a.png" align="center">
• RLax - 实现强化学习代理的库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/b8bfa3d3-948d-4c13-89dd-ebf6adc52750.png" align="center">
• JAX, M.D. - 加速、可微分的分子动力学。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/2a7b46eb-468a-4c0a-a32a-dface2e89208.png" align="center">
• Coax - 将强化学习论文转化为代码，轻松实现。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/ef4dd61c-080a-42a2-a433-60df8ef71bdc.png" align="center">
• Distrax - 重新实现 TensorFlow Probability，包含概率分布及其双射函数。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/c22b7056-5f10-4f40-bd41-2b1e6dc9e9b1.png" align="center">
• cvxpylayers - 构建可微分的凸优化层。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/ab6a7090-07a2-4f9d-817f-f96d9c1920bc.png" align="center">
• TensorLy - 简化张量学习。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/5012f55c-f2ef-4e56-bfc2-8c086d64381e.png" align="center">
• NetKet - 用于量子物理的机器学习工具箱。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/9190fd6a-3c02-40ca-9530-254aaa5d54dd.png" align="center">
• Fortuna - AWS 的深度学习不确定性量化库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/5125ed43-a286-45d7-ac57-042a0d519e55.png" align="center">
• BlackJAX - 针对 JAX 的采样器库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/f2f90696-51c6-4296-b558-6b1b56013b53.png" align="center">

<a name="new-libraries" />

新库

此部分包含制作精良且实用的库，但不一定经过大量用户严格测试。

• 神经网络库
• FedJAX - 基于Optax和Haiku在JAX中的联邦学习。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/4a0a0a60-98e9-4140-b9e7-675f76162f1c.png" align="center">
• Equivariant MLP - 构建等变神经网络层。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/b7d4acd4-16c6-4a59-ad53-76022bd60baf.png" align="center">
• jax-resnet - 实现和检查点在Flax中的ResNet变体。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/699d44c1-eadd-4c9c-997d-d00fb3fab435.png" align="center">
• Parallax - JAX中不可变的Torch模块。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/fd14d7ff-0c8d-4846-a398-a38f1f15f8e2.png" align="center">
• jax-unirep - 实现了用于蛋白质机器学习应用的UniRep模型的库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/c83eacb2-39aa-4713-a142-1db0ec287b10.png" align="center">
• jax-flows - JAX中的归一化流。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/d9992f72-4bf4-405a-87b9-4773209b43cd.png" align="center">
• sklearn-jax-kernels - 使用JAX的scikit-learn核矩阵。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/1a316264-0411-4ed3-b708-1365d86c5f92.png" align="center">
• jax-cosmo - 微分宇宙学库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/16342af1-9d17-4e1d-bba5-709dfb1ccab4.png" align="center">
• efax - JAX中的指数家族。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/7afaee1e-7c4c-45ba-8461-5758630660c3.png" align="center">
• mpi4jax - 在CPU和GPU上将MPI操作与Jax代码结合。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/c50f21df-ac9a-421d-bdf5-7f796959f28e.png" align="center">
• imax - 图像增强和变换。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/493d6196-efda-454c-818c-269c0d946add.png" align="center">
• FlaxVision - TorchVision的Flax版本。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/5bd87eca-7344-4a21-bac1-6046ce473308.png" align="center">
• Oryx - 基于程序转换的概率编程语言。
• Optimal Transport Tools - 解决优化传输问题的工具箱。
• delta PV - 具有自动微分功能的光伏模拟器。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/1e0df745-295f-450d-84c8-f5ded8162ee5.png" align="center">
• jaxlie - 刚体转换和优化的李理论库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/8df6b7a6-829f-48ca-b283-99903af52085.png" align="center">
• BRAX - 可微物理引擎，用于模拟环境以及训练这些环境的代理学习算法。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/bd6f75ac-40d4-4ff8-b9a6-23bc16de2281.png" align="center">
• flaxmodels - 适用于Jax/Flax的预训练模型。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/3b1f497d-91f9-4a1a-bc92-13638aed5f4e.png" align="center">
• CR.Sparse - 用于稀疏表示和压缩感知的XLA加速算法。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/6a56ce65-60aa-4c37-8646-234da81ecd60.png" align="center">
• exojax - 自动可微的系外行星/棕矮星光谱建模库，兼容JAX。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/76c59cc5-aa77-445b-9544-2909e7acab4a.png" align="center">
• JAXopt - JAX中的硬件加速(GPU/TPU)、可批处理和可微优化器。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/462322e7-9626-4846-9f2c-98c22c1155fc.png" align="center">
• PIX - JAX中的图像处理库，为JAX而生。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/d5c0a1c6-28ea-4370-a3cf-3b681873ba77.png" align="center">
• bayex - 由JAX驱动的贝叶斯优化库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/867a388b-53d3-471f-843c-4a38646dc9e4.png" align="center">
• JaxDF - 具有任意离散化的微分模拟框架。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/44108133-34dc-482f-9b32-ec59deac2872.png" align="center">
• tree-math - 将操作数组的函数转换为操作PyTrees的函数。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/2b7e3f6b-661f-4909-be01-99c62dd5223f.png" align="center">
• jax-models - 实现最初没有代码或用其他框架编写的研究论文。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/b231c3ea-2213-43b2-b56b-5288b07592be.png" align="center">
• PGMax - 一个用于构建离散概率图模型(PGM)并在JAX中进行推理的框架。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/eeabd421-06d8-4c15-bb6d-bd693c2e7304.png" align="center">
• EvoJAX - 硬件加速的神经进化库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/20c587f0-88e6-4e06-9cdf-6e42ab266cd5.png" align="center">
• evosax - 基于JAX的进化策略。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/16cc3f1d-166f-40e1-a559-b8692c58a6aa.png" align="center">
• SymJAX - 符号CPU/GPU/TPU编程。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/a6a4b10d-1630-4a91-a094-6974509aafd5.png" align="center">
• mcx - 表达和编译概率程序以实现高效推理。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/9a4413a6-f068-422f-b400-ffbe8875769f.png" align="center">
• Einshape - 基于DSL的JAX和其他框架的重构库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/19c5cbb7-d486-4ecb-a421-b0911c1a2afc.png" align="center">
• ALX - 使用交替最小二乘法进行分布式矩阵分解的开源库，更多信息见ALX: Large Scale Matrix Factorization on TPUs。
• Diffrax - JAX中的数值微分方程求解器。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/3093c6c8-5f0d-4231-8b7f-f3d65b06473f.png" align="center">
• tinygp - JAX中“最小”的高斯过程库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/8dd12f06-4bf6-48a9-85b9-9f2501761972.png" align="center">
• gymnax - 具有知名gym API的强化学习环境。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/f390a5ed-1e65-4ceb-9a08-a6311716d372.png" align="center">
• Mctx - 原生JAX中的蒙特卡洛树搜索算法。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/bbe7dd9d-d2e7-46bf-bb38-710889809759.png" align="center">
• KFAC-JAX - 用于神经网络的近似曲率二阶优化。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/083e8d5f-6389-485a-b516-acc2c966d689.png" align="center">
• TF2JAX - 将函数/图转换为JAX函数。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/15f494cb-9ae5-42b5-9fd8-942bb0442f3b.png" align="center">
• jwave - 用于可微分声学模拟的库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/77ee0789-a10d-4a31-971d-8332057f7228.png" align="center">
• GPJax - JAX中的高斯过程。
• Jumanji - 一套行业驱动的硬件加速强化学习环境，使用JAX编写。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/83ff6d70-4506-4bbc-a42e-e4dbaf3bb948.png" align="center">
• Eqxvision - Torchvision的Equinox版本。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/0969c5ee-4471-4eab-a168-ae665fb4d3d6.png" align="center">
• JAXFit - 用于非线性最小二乘问题的加速曲线拟合库（见arXiv论文）。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/58b19b95-8443-4898-b300-e55cc42bf2e6.png" align="center">
• econpizza - 使用JAX求解具有异质代理的宏观经济模型。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/f28f4d6b-cf90-43d3-a3a2-e949372ea4aa.png" align="center">
• SPU - 一个特定领域的编译器和运行时套件，用于使用JAX代码进行MPC（安全多方计算）。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/259d7ece-6191-4fb5-a5e9-97bd7e6449b8.png" align="center">
• jax-tqdm - 为JAX扫描和循环添加tqdm进度条。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/20a7661b-9f09-4612-9827-65946b0e215e.png" align="center">
• safejax - 使用🤗safetensors序列化JAX、Flax、Haiku或Objax模型参数。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/48c53067-23ce-4da4-8752-5109eb8379ec.png" align="center">
• Kernex - JAX中的可微分矩阵装饰器。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/857aadc1-9cb9-46e4-901b-ff1be3f08ae7.png" align="center">
• MaxText - 一个使用纯Python/JAX编写的简单、高性能和可扩展的JAX大语言模型，面向Google Cloud TPU。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/e742c83b-dac9-4c65-a661-cbd1b007c554.png" align="center">
• Pax - 一个基于JAX的大规模模型训练机器学习框架。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/30286269-2193-4e91-98ed-2ac639ceb39a.png" align="center">
• Praxis - Pax的层库，旨在供其他JAX基础的机器学习项目使用。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/8d2e659d-61f7-46b1-adb1-be6c9f1aea4f.png" align="center">
• purejaxrl - JAX中的可向量化端到端强化学习算法。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/535d1bea-bcf3-46e1-a850-0b0c4fdb4be8.png" align="center">
• Lorax - 自动将LoRA应用于JAX模型（Flax、Haiku等）。
• SCICO - JAX中的科学计算成像。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/1cd7f9cc-88b0-4f8f-a055-ceb2ae8b49c2.png" align="center">
• Spyx - 用于神经形态硬件机器学习的JAX尖峰神经网络。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/f8df508f-6d2f-4893-8623-804073876b24.png" align="center">
• BrainPy - 用Python编写的大脑动态编程。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/4ed7463b-da0e-4b89-b8db-c4b18f76fe5e.png" align="center">
• OTT-JAX - JAX中的最优传输工具。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/71d6540d-7e35-4307-bc12-5f1cada39c63.png" align="center">
• QDax - JAX中的质量多样性优化。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/f009d99c-ad4c-4a9e-b68a-c88885308006.png" align="center">
• JAX Toolbox - 使用T5x、Paxml和Transformer Engine等库在NVIDIA GPU上优化的JAX例子和夜间CI。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/76eaf61f-c247-479c-9b2a-4b225223b47b.png" align="center">
• Pgx - 带有AlphaZero示例的强化学习向量化棋盘游戏环境。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/32bda8a0-3096-4205-9974-b5286d5fa0e6.png" align="center">
• EasyDeL - EasyDeL 🔮是一个开源库，使您的训练更快、更优化，并提供用于JAX中训练和服务（Llama、MPT、Mixtral、Falcon等）的酷选项。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/8bcaef6b-18a2-4007-bae7-7cf569b90b95.png" align="center">
• XLB - Python中用于基于物理的机器学习的可微分大规模并行格子Boltzmann库。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/e7df89c6-b849-4359-b262-4744f27de8cc.png" align="center">
• dynamiqs - JAX中高性能和可微分的量子系统模拟。<img src="https://yellow-cdn.veclightyear.com/35dd4d3f/5b689e08-e9c3-426c-b46f-620002635f30.png" align="center"> <a name="models-and-projects" />

模型和项目

JAX

• Fourier Feature Networks - 官方实现 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains。
• kalman-jax - 使用迭代卡尔曼滤波和平滑进行马尔可夫（即时间）高斯过程的近似推理。
• jaxns - 在 JAX 中的嵌套采样。
• Amortized Bayesian Optimization - 与 Amortized Bayesian Optimization over Discrete Spaces 相关的代码。
• Accurate Quantized Training - 在 JAX 和 Flax 中运行和分析神经网络量化实验的工具和库。
• BNN-HMC - 为论文 What Are Bayesian Neural Network Posteriors Really Like? 提供的实现。
• JAX-DFT - 在 JAX 中的一维密度泛函理论 (DFT)，实现了 Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics。
• Robust Loss - 为论文 A General and Adaptive Robust Loss Function 提供的参考代码。
• Symbolic Functionals - 来自 Evolving symbolic density functionals 的演示。
• TriMap - TriMap: Large-scale Dimensionality Reduction Using Triplets 的官方 JAX 实现。

Flax

• Performer - Performer (通过 FAVOR+ 的线性 Transformer) 架构的 Flax 实现。
• JaxNeRF - 支持多设备 GPU/TPU 的 NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis 实现。
• mip-NeRF - Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields 的官方实现。
• RegNeRF - RegNeRF: Regularizing Neural Radiance Fields for View Synthesis from Sparse Inputs 的官方实现。
• Big Transfer (BiT) - Big Transfer (BiT): General Visual Representation Learning 的实现。
• JAX RL - 强化学习算法的实现。
• gMLP - Pay Attention to MLPs 的实现。
• MLP Mixer - MLP-Mixer: An all-MLP Architecture for Vision 的简洁实现。
• Distributed Shampoo - Second Order Optimization Made Practical 的实现。
• NesT - Aggregating Nested Transformers 的官方实现。
• XMC-GAN - Cross-Modal Contrastive Learning for Text-to-Image Generation 的官方实现。
• FNet - FNet: Mixing Tokens with Fourier Transforms 的官方实现。
• GFSA - Learning Graph Structure With A Finite-State Automaton Layer 的官方实现。
• IPA-GNN - Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks 的官方实现。
• Flax Models - 在 Flax 中实现的模型和方法集合。
• Protein LM - 实现了 BERT 和自回归蛋白质模型，如 Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences 和 ProGen: Language Modeling for Protein Generation 中所述。
• Slot Attention - Differentiable Patch Selection for Image Recognition 的参考实现。
• Vision Transformer - An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale 的官方实现。
• FID computation - 将 mseitzer/pytorch-fid 移植到 Flax 上。
• ARDM - Autoregressive Diffusion Models 的官方实现。
• D3PM - Structured Denoising Diffusion Models in Discrete State-Spaces 的官方实现。
• Gumbel-max Causal Mechanisms - Learning Generalized Gumbel-max Causal Mechanisms 的代码，额外代码见 GuyLor/gumbel_max_causal_gadgets_part2。
• Latent Programmer - ICML 2021 论文 Latent Programmer: Discrete Latent Codes for Program Synthesis 的代码。
• SNeRG - ICLR 2022 论文 Baking Neural Radiance Fields for Real-Time View Synthesis 的官方实现。
• Spin-weighted Spherical CNNs - Spin-Weighted Spherical CNNs 的改编。
• VDVAE - Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images 的改编，原始代码见 openai/vdvae。
• MUSIQ - ICCV 2021 论文 MUSIQ: Multi-scale Image Quality Transformer 的检查点和模型推理代码。
• AQuaDem - Continuous Control with Action Quantization from Demonstrations 的官方实现。
• Combiner - Combiner: Full Attention Transformer with Sparse Computation Cost 的官方实现。
• Dreamfields - ICLR 2022 论文 Progressive Distillation for Fast Sampling of Diffusion Models 的官方实现。
• GIFT - Gradual Domain Adaptation in the Wild: When Intermediate Distributions are Absent 的官方实现。
• Light Field Neural Rendering - Light Field Neural Rendering 的官方实现。
• Sharpened Cosine Similarity in JAX by Raphael Pisoni - Sharpened Cosine Similarity 层的 JAX/Flax 实现。
• GNNs for Solving Combinatorial Optimization Problems - 使用 JAX + Flax 实现的 Combinatorial Optimization with Physics-Inspired Graph Neural Networks。

Haiku

• AlphaFold - AlphaFold v2.0 推理流程的实现，见论文 Highly accurate protein structure prediction with AlphaFold。
• Adversarial Robustness - 参考代码，见论文 Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples 和 Fixing Data Augmentation to Improve Adversarial Robustness。
• Bootstrap Your Own Latent - 自监督学习新方法，见论文 Bootstrap your own latent: A new approach to self-supervised Learning。
• Gated Linear Networks - GLN 是一类不依赖反向传播的神经网络。
• Glassy Dynamics - 开源实现，见论文 Unveiling the predictive power of static structure in glassy systems。
• MMV - 实现模型代码，见论文 Self-Supervised MultiModal Versatile Networks。
• Normalizer-Free Networks - 官方 Haiku 实现，见论文 NFNets。
• NuX - 使用 JAX 的正则化流。
• OGB-LSC - 包含 DeepMind 在 OGB Large-Scale Challenge (OGB-LSC) 中的 PCQM4M-LSC (量子化学) 和 MAG240M-LSC (学术图) 赛道的参赛代码。
• Persistent Evolution Strategies - 代码实现，见论文 Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies。
• Two Player Auction Learning - JAX 实现，见论文 Auction learning as a two-player game。
• WikiGraphs - 基准代码，见论文 WikiGraphs: A Wikipedia Text - Knowledge Graph Paired Datase。

Trax

• Reformer - 高效 Transformer 架构 Reformer 的实现。

NumPyro

• lqg - 线性-二次高斯问题的贝叶斯逆最优控制的官方实现，见论文 Putting perception into action with inverse optimal control for continuous psychophysics。

<a name="videos" />

视频

• NeurIPS 2020: JAX Ecosystem Meetup - JAX 在 DeepMind 的使用及工程师、科学家和 JAX 核心团队之间的讨论。
• Introduction to JAX - 从零开始用 JAX 构建简单神经网络。
• JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas - JAX 的核心设计、推动新研究的方式及如何使用它。
• Bayesian Programming with JAX + NumPyro — Andy Kitchen - 使用 NumPyro 进行贝叶斯建模的介绍。
• JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne - 在 Program Transformations for Machine Learning 研讨会上，JAX 简介演讲。
• JAX on Cloud TPUs | NeurIPS 2020 | Skye Wanderman-Milne and James Bradbury - TPU 主机访问演示。
• Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020 - 教程由 Zico Kolter、David Duvenaud 和 Matt Johnson 创建，Colab 笔记本见 Deep Implicit Layers。
• Solving y=mx+b with Jax on a TPU Pod slice - Mat Kelcey - 一系列四部分 YouTube 教程及 Colab 笔记，本系列从 Jax 基础讲起，最后讨论在 v3-32 TPU Pod 切片上进行数据并行训练。
• JAX, Flax & Transformers 🤗 - 围绕 JAX / Flax、Transformers、大规模语言建模及其他精彩话题的三天讲座。

<a name="papers" />

论文

此部分包含围绕 JAX 的论文（例如基于 JAX 的库白皮书、JAX 研究等）。使用 JAX 实现的论文见 Models/Projects 章节。

<!--lint ignore awesome-list-item-->

• Compiling machine learning programs via high-level tracing. Roy Frostig, Matthew James Johnson, Chris Leary. MLSys 2018. - 描述 JAX 早期版本的白皮书，详述了计算如何被追踪和编译。
• JAX, M.D.: A Framework for Differentiable Physics. Samuel S. Schoenholz, Ekin D. Cubuk. NeurIPS 2020. - 介绍 JAX, M.D.，一个包含仿真环境、相互作用势、神经网络等在内的可微物理学库。
• Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization. Pranav Subramani, Nicholas Vadivelu, Gautam Kamath. arXiv 2020. - 使用 JAX 的 JIT 和 VMAP 实现比现有库更快的差分隐私 SGD。
• XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python. Mohammadmehdi Ataei, Hesam Salehipour. arXiv 2023. - 描述 XLB 库的白皮书：基准测试、验证及库的更多细节。

<!--lint enable awesome-list-item--> <a name="tutorials-and-blog-posts" />

教程和博文

• 使用 JAX 加速我们的研究 by David Budden 和 Matteo Hessel - 描述了 DeepMind 中 JAX 和 JAX 生态系统的现状。
• 开始使用 JAX (MLPs, CNNs & RNNs) by Robert Lange - 从头开始使用基本的 JAX 算子构建神经网络模块。
• 学习 JAX: 从线性回归到神经网络 by Rito Ghosh - 一个温和的介绍，使用 JAX 实现线性和逻辑回归，以及神经网络模型，并使用它们解决实际问题。
• 教程：使用 JAX 和 Flax Linen 进行图像分类 by 8bitmp3 - 学习如何使用 Flax 的 Linen API 创建一个简单的卷积网络，并训练其识别手写数字。
• 使用 JAX by Nick Doiron - 比较了 Flax、Haiku 和 Objax 在 Kaggle 花卉分类挑战中的表现。
• 在 JAX 中的 50 行元学习 by Eric Jang - 介绍了 JAX 和元学习。
• 在 JAX 中的 100 行归一化流 by Eric Jang - 简明实现 RealNVP。
• 在 GPU/TPU 上进行可微路径追踪 by Eric Jang - 实现路径追踪的教程。
• 集成网络 by Mat Kelcey - 集成网络是一种将多个模型集成在一个逻辑模型中的方法。
• 分布外 (OOD) 检测 by Mat Kelcey - 实现不同的 OOD 检测方法。
• 使用 JAX 理解自动微分 by Srihari Radhakrishna - 使用 JAX 理解自动微分的工作原理。
• 从 PyTorch 到 JAX: 向净化状态代码的神经网络框架过渡 by Sabrina J. Mielke - 展示如何从类似 PyTorch 的编码风格转向更具功能性的编码风格。
• 使用自定义 C++ 和 CUDA 代码扩展 JAX by Dan Foreman-Mackey - 介绍为 JAX 提供自定义操作所需的基础设施的教程。
• 在 JAX 中进化神经网络 by Robert Tjarko Lange - 探索如何利用 JAX 推动下一代可扩展的神经进化算法。
• 使用 JAX 探索超参数元损失图景 by Luke Metz - 演示如何使用 JAX 进行内损失优化（使用 SGD 和 Momentum）、外损失优化（使用梯度）以及外损失优化（使用进化策略）。
• 在 JAX 中实现确定性 ADVI by Martin Ingram - 使用 JAX 轻松、清晰地实现自动微分变分推断（ADVI）的教程。
• 进化通道选择 by Mat Kelcey - 训练一个能适应不同分辨率下输入通道组合的分类模型，然后使用遗传算法决定最佳组合。
• JAX 入门 by Kevin Murphy - 介绍语言各方面并应用于简单的机器学习问题的 Colab。
• 在 JAX 中编写 MCMC 采样器 by Jeremie Coullon - 关于在 JAX 中编写 MCMC 采样器的不同方法和速度基准的教程。
• 如何为 JAX 扫描和循环添加进度条 by Jeremie Coullon - 教程，教你如何使用 host_callback 模块为编译的循环添加进度条。
• 开始使用 JAX by Aleksa Gordić - 一系列讲解如何从零学习 JAX 知识到在 Haiku 中构建神经网络的笔记本和视频。
• 在 JAX+FLAX 中编写训练循环 by Saurav Maheshkar 和 Soumik Rakshit - 关于在 JAX、Flax 和 Optax 中编写简单的端到端训练和评估管道的教程。
• 在 JAX 中实现 NeRF by Soumik Rakshit 和 Saurav Maheshkar - 关于 JAX 中进行由神经辐射场表示的场景的 3D 体素渲染教程。
• 使用 JAX+Flax 进行深度学习教程 by Phillip Lippe - 一系列讲解各种深度学习概念的笔记，从基础（例如 JAX/Flax 简介、激活函数）到最新进展（例如视觉变压器、SimCLR），并翻译成 PyTorch 代码。
• 使用 PureJaxRL 实现 4000 倍加速 - 介绍 JAX 如何通过向量化大幅加速强化学习训练的一篇博客文章。

书籍

• Jax 实战 - 一本关于使用 JAX 进行深度学习和其他数学密集型应用的实践指南。

社区

• JaxLLM（非官方）Discord
• JAX GitHub 讨论
• Reddit

贡献

欢迎贡献！请先阅读 贡献指南。