Oryx 是一个基于 JAX 的概率编程和深度学习库。

import oryx
import jax.numpy as jnp
ppl = oryx.core.ppl
tfd = oryx.distributions

# Define sampling function
def sample(key):
  x = ppl.random_variable(tfd.Normal(0., 1.))(key)
  return jnp.exp(x / 2.) + 2.

# Transform sampling function into a log-density function
ppl.log_prob(sample)(1.)  # ==> -0.9189
Oryx 的方法是公开一组会组合和集成 JAX 现有转换的函数转换。如需安装 Oryx,您可以运行以下命令:
 pip install --upgrade oryx