Levanter是一个基于JAX和Haliax的框架,用于训练大型语言模型和其他基础模型,专注于可读性、可扩展性和可复现性。
EasyLM是一个基于JAX/Flax的简单可扩展大语言模型训练框架,提供预训练、微调、评估和服务的一站式解决方案。
Marin是一个开源的基础模型研究和开发框架,专注于语言模型的训练和开发,强调实验的可重现性。
PyMC是一个用于Python的概率编程库,允许用户使用简单的Python API构建贝叶斯模型,并使用马尔可夫链蒙特卡洛(MCMC)方法进行拟合。
TensorFlow Probability是一个基于TensorFlow构建的Python库,用于概率推理和统计分析,支持在现代硬件上结合概率模型和深度学习。
Distrax是Google DeepMind开发的轻量级概率分布和双射器库,作为TensorFlow Probability的JAX原生实现,专注于可扩展性和易用性。
JAX MD是一个基于JAX的分子动力学模拟工具包,专门用于处理粒子集合的物理模拟和机器学习应用
Brax是Google开发的一个快速且完全可微分的物理引擎,专为加速硬件设计,用于机器人、强化学习等仿真密集型应用研究。
Grain是一个用于读取和处理机器学习训练数据的Python库,专门为JAX模型设计,具有灵活、快速和确定性的特点。
🤗 Datasets是一个用于轻松访问和共享AI数据集的库,支持音频、计算机视觉和自然语言处理任务。
Orbax是Google开发的JAX训练工具集,提供检查点管理和模型导出功能,专为多主机、多设备环境设计。
Blackjax是一个基于JAX的采样器库,提供快速、稳健且经过充分测试的采样算法,适用于CPU和GPU计算。
NumPyro是基于NumPy的概率编程库,提供灵活的贝叶斯建模和推理功能,支持多种机器学习算法。
Flax是一个用于JAX的神经网络库,为研究人员和开发者提供灵活的用户体验,使JAX在神经网络应用中的全部能力得以发挥。
Equinox是一个基于JAX的深度学习库,提供PyTorch风格的神经网络构建语法和丰富的PyTree操作功能,兼容整个JAX生态系统。
Optax是一个用于JAX的梯度处理和优化库,为优化参数模型(如深度神经网络)提供可组合的构建块。
Optimistix是一个基于JAX和Equinox的非线性优化库,提供根查找、最小化、固定点和最小二乘等求解器。
Lineax是一个JAX库,用于线性求解和线性最小二乘法,支持PyTree值矩阵和向量操作。
Diffrax是一个基于JAX的数值微分方程求解器库,提供ODE/SDE/CDE等多种微分方程求解功能。
Bokeh是一个Python库,用于为现代Web浏览器创建交互式可视化图表,无需编写JavaScript即可构建从简单图表到复杂仪表板的可视化应用。