Effect Handling for Composable Program Transformations in Edward2

TL;DR

This paper explores the use of algebraic effect handlers as a modular abstraction for implementing probabilistic programming languages, enabling flexible and composable model transformations and inference algorithms.

Contribution

It interprets Edward2's design as an effect-handling mechanism and extends it to support nested, composable transformations for probabilistic programming.

Findings

Enables straightforward implementation of complex model transformations

Supports nested, composable probabilistic effects

Improves modularity and flexibility in probabilistic programming systems

Abstract

Algebraic effects and handlers have emerged in the programming languages community as a convenient, modular abstraction for controlling computational effects. They have found several applications including concurrent programming, meta programming, and more recently, probabilistic programming, as part of Pyro's Poutines library. We investigate the use of effect handlers as a lightweight abstraction for implementing probabilistic programming languages (PPLs). We interpret the existing design of Edward2 as an accidental implementation of an effect-handling mechanism, and extend that design to support nested, composable transformations. We demonstrate that this enables straightforward implementation of sophisticated model transformations and inference algorithms.