Skip to content

Tutorial notebook. Source .ipynb on GitHub.

ProbJAX PPL Basics¶

This notebook uses the probabilistic-first API names from probjax.core: observe, do, and log_joint_fn (with backward-compatible aliases still available).

import jax
import jax.numpy as jnp

from probjax.core import do, joint_sample, log_joint_fn, observe, scope, substitute, trace
from probjax.core.custom_primitives.random_variable import rv_p
from probjax.stats import norm

def model(key):
    k1, k2 = jax.random.split(key)
    z = rv_p.bind(k1, 0.0, 1.0, dist=norm, name="z")
    y = rv_p.bind(k2, z, 0.5, dist=norm, name="y")
    return y

1) Joint sampling¶

joint_sample(model)(jax.random.PRNGKey(0))

{'z': Array(1.0040143, dtype=float32), 'y': Array(-0.21721363, dtype=float32)}

2) Observe data and evaluate log-joint¶

obs_y = jnp.array(0.25)
observed_model = observe(model, {"y": obs_y})
latent_samples = joint_sample(observed_model)(jax.random.PRNGKey(1))
log_joint_fn(observed_model)(z=latent_samples["z"])

Array(-1.4451859, dtype=float32)

3) Substitute in `condition` mode vs `do` intervention¶

fixed_z = jnp.array(-0.4)

conditioned_substitute = substitute(model, {"z": fixed_z}, mode="condition")
conditioned_samples = joint_sample(conditioned_substitute)(jax.random.PRNGKey(2))
conditioned_log_joint = log_joint_fn(conditioned_substitute)(y=conditioned_samples["y"])

do_model = do(model, {"z": fixed_z})
do_samples = joint_sample(do_model)(jax.random.PRNGKey(3))
do_log_joint = log_joint_fn(do_model)(y=do_samples["y"])

{
    "conditioned_samples": conditioned_samples,
    "conditioned_log_joint": conditioned_log_joint,
    "do_samples": do_samples,
    "do_log_joint": do_log_joint,
}

{'conditioned_samples': {'y': Array(0.91310763, dtype=float32)},
 'conditioned_log_joint': Array(-4.673233, dtype=float32),
 'do_samples': {'y': Array(-1.3716108, dtype=float32)},
 'do_log_joint': Array(-2.1138465, dtype=float32)}

mode="condition" keeps the substituted site's density term. do(...) removes the intervened site's own density term.

4) Site-aware tracing¶

trace(model, sites=True)(jax.random.PRNGKey(4))

{'z': {'name': 'z',
  'kind': 'sample',
  'kind_probabilistic': 'sample',
  'kind_legacy': 'sample',
  'value': Array(0.11766523, dtype=float32),
  'log_prob': Array(-0.92586106, dtype=float32),
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()},
 'y': {'name': 'y',
  'kind': 'sample',
  'kind_probabilistic': 'sample',
  'kind_legacy': 'sample',
  'value': Array(0.29400343, dtype=float32),
  'log_prob': Array(-0.2879817, dtype=float32),
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()}}

trace(conditioned_substitute, sites=True)(jax.random.PRNGKey(5))

{'z': {'name': 'z',
  'kind': 'condition',
  'kind_probabilistic': 'condition',
  'kind_legacy': 'replay',
  'value': Array(-0.4, dtype=float32, weak_type=True),
  'log_prob': Array(-0.9989385, dtype=float32, weak_type=True),
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()},
 'y': {'name': 'y',
  'kind': 'sample',
  'kind_probabilistic': 'sample',
  'kind_legacy': 'sample',
  'value': Array(-1.0899351, dtype=float32),
  'log_prob': Array(-1.1778122, dtype=float32),
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()}}

trace(do_model, sites=True)(jax.random.PRNGKey(6))

{'z': {'name': 'z',
  'kind': 'do',
  'kind_probabilistic': 'do',
  'kind_legacy': 'intervene',
  'value': Array(-0.4, dtype=float32, weak_type=True),
  'log_prob': None,
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()},
 'y': {'name': 'y',
  'kind': 'sample',
  'kind_probabilistic': 'sample',
  'kind_legacy': 'sample',
  'value': Array(-0.8336562, dtype=float32),
  'log_prob': Array(-0.6019068, dtype=float32),
  'dist': <probjax.stats.continuous.norm.norm_gen at 0x10ec87e00>,
  'shape': ()}}

5) Strict vs partial log-joint evaluation¶

strict_log_joint = log_joint_fn(model)
try:
    strict_log_joint(z=jnp.array(0.1))
except KeyError as err:
    print("strict mode error:", err)

strict mode error: "Missing joint sample for random variable 'y'."

partial_log_joint = log_joint_fn(model, allow_partial=True)
partial_log_joint(z=jnp.array(0.1))

Array(-0.9239385, dtype=float32)

6) Scoped auto-generated site names¶

def scoped_model(key):
    with scope("outer"):
        z = rv_p.bind(key, 0.0, 1.0, dist=norm)
    return z

joint_sample(scoped_model)(jax.random.PRNGKey(7))

{'outer__norm_0': Array(0.45123515, dtype=float32)}