JaxDecompiler 是一个将任何 JAX 函数转换为可编辑的 Python 代码的工具，特别适用于编辑由梯度函数生成的 JAX 函数，并在逆向工程、理解、定制化和软件互通方面简化了过程。它在深度学习和梯度驱动软件方面具有实际应用，我们通过演示证明了反编译代码的速度性能与原始代码相似。

Mar, 2024

我们通过向 JAX 增加自动区分高阶函数（functional 和 operators）的能力，将函数表示为数组的一种泛化，采用 JAX 现有的原始系统来实现高阶函数。我们提供了一组原始操作符，作为构建几种关键类型的 functional 的基础构建块。对于每个引入的原始操作符，我们推导和实现了线性化和转置规则，与 JAX 内部的前向和反向模式自动区分协议保持一致。这种增强功能允许在传统用于函数的语法中进行 functional 区分。得到的 functional 梯度本身就是可以在 Python 中调用的函数。我们通过一些应用展示了这个工具的效能和简洁性，这些应用中 functional 导数是不可或缺的。本研究的源代码可在此 https URL 中获取。

Nov, 2023

BlackJAX 是一个库，用于实现在贝叶斯计算中常用的采样和变分推断算法。它采用了功能性的方法来实现算法，以便于使用、速度快以及具备模块化。

Feb, 2024

介绍了 Equinox，它是一种小型神经网络库，采用类似 PyTorch 的基于类的方法，同时使用 JAX-like 的函数编程，它通过 PyTree 和函数变换解决了 OO/functional 的差异而不需要引入新的程序抽象。

Oct, 2021

FAX 是一个基于 JAX 的库，旨在支持数据中心和跨设备应用中的大规模分布式和联合计算。FAX 利用 JAX 的分片机制，实现了针对 TPUs 和 JAX 运行时的本地化支持。FAX 将联合计算的构建块嵌入 JAX 的原语中，并提供了对 XLA HLO 的翻译，联合自动微分功能的完整实现，以及对现有生产跨设备联合计算系统的解释。我们展示了 FAX 在数据中心中提供了易于编程、高性能和可扩展的联合计算框架。

Mar, 2024

本文介绍了一种新的 Python 软件包 BindsNET，用于模拟脉冲神经网络，专门针对机器学习和强化学习，旨在促进脉冲神经网络在大规模机器学习实验中的使用，提供了用户友好、简洁的语法。

Jun, 2018

提出 Binder，一种基于神经符号方法的框架，采用 GPT-3 Codex 作为输入的语言模型，通过少量标注建立程序，实现对丰富的问题进行回答，取得了在 WikiTableQuestions 和 TabFact 两个数据集上的最先进结果。

Oct, 2022

FedJAX 是一个基于 JAX 的开放源代码库，用于联邦学习模拟，旨在使开发和评估联邦算法更快、更容易，并强调隐私和安全。FedJAX 可以在几分钟内使用 TPU 训练 EMNIST 数据集上的模型，使用标准超参数在约一小时内训练 Stack Overflow 数据集。

Aug, 2021

SynJax 提供了对结构化分布推理算法的高效向量化实现，包括对齐、标记、分段、句法树和生成树，通过使用 SynJax，我们能构建大规模可微分的模型，能明确地对数据中的结构进行建模。

Aug, 2023

本论文介绍了一种基于 JAX 的裁剪和稀疏训练库 JaxPruner，旨在通过提供流行的裁剪和稀疏训练算法的简洁实现，减少内存和延迟开销，从而加速稀疏神经网络的研究，此外，JaxPruner 的算法使用公共 API，并与流行的优化库 Optax 无缝配合，使其易于与现有的基于 JAX 的库集成。本论文提供了四个不同代码库的示例，并在流行的基准测试上提供了基准实验。

Apr, 2023