NumPyro 中用于灵活且加速的概率编程的可组合效应
机器学习
2019-12-30 v1 人工智能
机器学习
编程语言
摘要
NumPyro 是一个轻量级库,为 Pyro 概率编程语言提供替代的 NumPy 后端,具有相同的建模接口、语言原语和效应处理抽象。效应处理器使得 Pyro 的建模 API 能够扩展到 NumPyro,尽管其构建于基于 JAX 的根本不同的函数式后端之上。在本工作中,我们展示了将 Pyro 的效应处理器与 JAX 中实现硬件加速、自动微分和向量化的程序变换相组合的能力。特别地,NumPyro 提供了 No-U-Turn 采样器(NUTS)的迭代公式,可进行端到端 JIT 编译,从而在小数据集和大数据集两种情形下均比现有替代方案快得多。
引用
@article{arxiv.1912.11554,
title = {Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro},
author = {Du Phan and Neeraj Pradhan and Martin Jankowiak},
journal= {arXiv preprint arXiv:1912.11554},
year = {2019}
}
备注
10 pages, 2 figures; NeurIPS 2019 Program Transformations for Machine Learning Workshop