Implementing a Simple Probabilistic Programming Language (PPL) in Python

The article introduces probabilistic programming languages (PPLs) and demonstrates how to implement a simple PPL using Python, targeting statisticians, AI researchers, and curious programmers familiar with Bayesian inference.

Related research notes that few Python‑based PPLs exist, citing works such as "The Design and Implementation of Probabilistic Programming Languages" (JavaScript implementation) and the PyMC3 developer guide.

Implementation – High‑level representation shows the core model API:

mu = LatentVariable("mu", Normal, [0.0, 5.0])
y_bar = ObservedVariable("y_bar", Normal, [mu, 1.0], observed=3.0)
evaluate_log_density(y_bar, {"mu": 4.0})

The two statements define a joint probability distribution whose PDF is visualized with both a probabilistic graphical model and a directed factor graph.

Distribution class is defined using SciPy’s norm :

from scipy.stats import norm

class Distribution:
    @staticmethod
    def log_density(point, params):
        raise NotImplementedError("Must be implemented by a subclass")

class Normal(Distribution):
    @staticmethod
    def log_density(point, params):
        return float(norm.logpdf(point, params[0], params[1]))

Variables and DAG are represented by LatentVariable and ObservedVariable classes, each storing a name, distribution class, arguments, and (for observed variables) a concrete value. The DAG is built from dist_args , which may be constants, latent variables, or observed variables.

class LatentVariable:
    def __init__(self, name, dist_class, dist_args):
        self.name = name
        self.dist_class = dist_class
        self.dist_args = dist_args

class ObservedVariable:
    def __init__(self, name, dist_class, dist_args, observed):
        self.name = name
        self.dist_class = dist_class
        self.dist_args = dist_args
        self.observed = observed

Log‑density evaluation traverses the DAG with a depth‑first search, collects variables, resolves each argument to a numeric value (using a latent_values dictionary for latent variables), and accumulates the log‑density:

def evaluate_log_density(variable, latent_values):
    visited = set()
    variables = []
    def collect_variables(v):
        if isinstance(v, float):
            return
        visited.add(v)
        variables.append(v)
        for arg in v.dist_args:
            if arg not in visited:
                collect_variables(arg)
    collect_variables(variable)
    log_density = 0.0
    for var in variables:
        dist_params = []
        for arg in var.dist_args:
            if isinstance(arg, float):
                dist_params.append(arg)
            elif isinstance(arg, LatentVariable):
                dist_params.append(latent_values[arg.name])
        if isinstance(var, LatentVariable):
            log_density += var.dist_class.log_density(latent_values[var.name], dist_params)
        elif isinstance(var, ObservedVariable):
            log_density += var.dist_class.log_density(var.observed, dist_params)
    return log_density

An example verifies the implementation:

mu = LatentVariable("mu", Normal, [0.0, 5.0])
y_bar = ObservedVariable("y_bar", Normal, [mu, 1.0], observed=5.0)
latent_values = {"mu": 4.0}
print(evaluate_log_density(y_bar, latent_values))  # => -4.267314978843446

Conclusion and future work summarises that distributions, variable DAGs, and log‑density computation are core components of a PPL, and suggests extensions such as tensor support, hierarchical models, prior‑predictive sampling, and integration with computational‑graph frameworks like JAX or TensorFlow for performance and automatic differentiation.

Posterior grid approximation (outside the main article) demonstrates using the log‑density function to evaluate a grid of mu values and plot the posterior mode, illustrating how the PPL can be used for simple Bayesian inference.

import numpy as np, pandas as pd, altair as alt
from smolppl import Normal, LatentVariable, ObservedVariable, evaluate_log_density
mu = LatentVariable("mu", Normal, [0.0, 5.0])
y_bar = ObservedVariable("y_bar", Normal, [mu, 1.0], observed=1.5)
grid = np.linspace(-4, 4, 20)
evals = [evaluate_log_density(y_bar, {"mu": m}) for m in grid]
chart = alt.Chart(pd.DataFrame({"mu": grid, "logdensity": evals})).mark_line(point=True).encode(x='mu', y='logdensity')
chart.show()