Edward2: A Low-Level Probabilistic Programming Layer for TensorFlow, JAX and NumPy

The problem Edward2 solves: models that are programs, not graphs

Most probabilistic modelling libraries ask you to declare a model in a form the library understands. Edward2 asks for a Python function. The README describes the project as a simple probabilistic programming language whose purpose is to let you write models as probabilistic programs and then manipulate a model's computation for flexible training and inference. That second half is the actual point. Writing the model is easy in many libraries. Being able to reach into the computation and change how it runs, without rewriting the model, is the part Edward2 is built around. The audience follows from that. This is not a tool for someone who wants to fit a Bayesian regression by calling one function. It is for someone implementing a variational objective, a custom posterior approximation, or a training loop that needs to swap sampling for something else. The README is explicit that the core utilities are fairly low-level, and it redirects readers who want a high-level module for uncertainty modelling to Bayesian Layers. That redirection is the clearest statement of scope in the whole document.

Random variables as the unit of model structure

A model in Edward2 is built from RandomVariables. According to the README, a random variable carries a probability distribution accessible as rv.distribution, which is a TensorFlow Distribution instance governing log_prob and sample. Construction mirrors TensorFlow Distributions: ed.Normal(loc=0., scale=1.) returns a RandomVariable, and the README shows normal_rv.distribution.log_prob(1.231) returning a tensor. By default, instantiating a random variable creates a sampling op so that rv.value equals rv.distribution.sample(). The default number of samples is one, controllable through the sample_shape argument to rv, and if you pass an optional value argument no sampling op is created at all. That last detail matters more than it looks. Supplying value is how you inject observed data or a variational sample into a program without changing the program's text. Random variables also interoperate with TensorFlow ops, which operate on the sample: the README shows x + y, x / y, tf.tanh(x * y) and indexing x[1] all producing tensors or further random variables. A probabilistic model is then just a function that instantiates one or more of these. The README's logistic_regression example takes a features tensor, draws coeffs and intercept from Normal priors, and returns an ed.Bernoulli over the logits. Inputs to the function are values the model conditions on.

Tracing: the mechanism that makes inference possible

If random variables were the whole story, Edward2 would be a thin wrapper over TensorFlow Distributions. Tracing is what separates it. The README defines a tracer as a function that acts on another function f and its arguments, performing computations before returning an output, typically f(*args, **kwargs). The ed.trace context manager pushes tracers onto a stack, and any traceable function is intercepted by the stack. The README states that all random variable constructors are traceable, which is the hook that makes the pattern work. The consequence is that you can call the same logistic_regression function twice and get different behaviour, because a tracer on the stack can intercept each ed.Normal and ed.Bernoulli construction as it happens. That is how you would substitute a posterior for a prior, or collect log probabilities, without editing the model. The README also shows a second pattern: writing a separate function, logistic_regression_posterior, that returns ed.MultivariateNormalTriL and ed.Normal random variables parameterised by tf.Variable objects. Here the variational parameters are ordinary TensorFlow variables, and the posterior is a program in the same sense the prior is. Note the manual constraint handling: the example uses tril_with_diag_softplus_and_shift for the scale_tril argument and tf.nn.softplus(intercept_scale) + 1e-5 for the intercept scale, so positivity is the modeller's job.

Installation, backends and the version you actually get

The README recommends the development version: pip install "edward2 @ git+https://github.com/google/edward2.git". A stable release exists via pip install edward2, but the README is candid that it is rarely updated, describing the project as a passion project maintained by part-timers where scheduling releases sucks up time. Treat that as a design constraint, not a complaint. If you install from PyPI you are probably getting an old snapshot, and the README's own recommendation points at git. Installing edward2 does not install any backend. You add one explicitly, for example pip install edward2[tensorflow], replacing tensorflow with the appropriate backend. TensorFlow is the default; JAX and NumPy are supported, and the README says to see below to activate them, though the supplied material does not include that section, so the exact activation mechanism cannot be confirmed here. There is also an extra for TensorFlow's nightly package, pip install edward2[tf-nightly], for cases where Edward2 uses the latest TensorFlow changes. The repository is organised into edward2/ for library code, examples/, and experimental/ for what the README calls active research projects. That third directory is worth reading before you import anything from it.

Where Edward2 is the wrong tool

The clearest limitation is stated by the project itself: the core utilities are low-level. There is no high-level uncertainty module in the core, and the README sends you to Bayesian Layers for that and to Uncertainty Baselines if you want research-ready code to build on. If your task is to add predictive uncertainty to an existing network with minimal code, Edward2's random variables and tracer stack are machinery you would have to assemble yourself. A second limitation is release cadence. The README says the stable version is very rarely updated, so bug fixes and backend compatibility changes land on main rather than in a versioned artifact. That means pinning a commit hash, and it means an upgrade is a deliberate act rather than a version bump. A third is backend coupling. The default path is TensorFlow, and the README notes that Edward2 sometimes depends on the latest TensorFlow changes, which is why the tf-nightly extra exists. A project that needs a frozen, reproducible TensorFlow version may find itself fighting that. Finally, the experimental/ directory is described as holding active research projects. Research code has no compatibility promise, and the README does not offer one.

Edward2 against Pyro and NumPyro

The natural comparison is Pyro or NumPyro, which are also probabilistic programming languages with effect-handling style interception. The difference is in what the model is written against. In Edward2, random variables are TensorFlow Distribution instances, so the distribution library, its parameterisations and its constraints are the ones you already use in TensorFlow. The README's posterior example is a good illustration: scale_tril is built with tfp.trainable_distributions.tril_with_diag_softplus_and_shift, and the intercept scale uses tf.nn.softplus. You are inside the TensorFlow and TensorFlow Probability ecosystem, including its optimisers and tf.Variable semantics. Pyro and NumPyro instead bring their own distribution library and their own inference primitives, with NumPyro running on JAX rather than TensorFlow. The practical consequence is where your existing code lives. If your training loop, data pipeline and deployment target are TensorFlow, Edward2's tracing sits on top of that stack rather than beside it. If they are not, you are adopting TensorFlow to get Edward2. Edward2 also offers JAX and NumPy backends, but the supplied README does not describe how those backends change the tracing or distribution semantics, so anyone whose reason for choosing Edward2 is the JAX backend should read that section of the repository before deciding.

Maintenance, licensing and what to check before adopting

Edward2 is Apache-2.0 licensed, which permits commercial use and modification and includes an explicit patent grant. That is a permissive licence, and it is the same licence used across much of the TensorFlow ecosystem, so combining Edward2 with TensorFlow code raises no obvious licensing friction. This is not legal advice; read the licence text if your organisation has specific requirements. On maintenance, the README's own framing is the most useful signal: part-time maintainers, a stable release that is rarely refreshed, and an experimental directory that is explicitly research. Plan for installation from git at a pinned commit, and plan for the possibility that a TensorFlow upgrade forces you onto tf-nightly. There are no release notes in the supplied material, so there is no basis for judging how often the library code changes or how breaking changes are communicated. Before adopting, verify three things against the repository rather than this article: the JAX and NumPy backend activation instructions, the tracer implementation in edward2/tracer.py, and whether anything you need lives under experimental/. If the answer to the third is yes, you are depending on code the README does not present as supported.

Editorial conclusion

Adopt Edward2 if you are writing custom inference and want the model to be an ordinary Python function you can intercept, and if you can pin a git commit because the maintainers state they very rarely update the stable release. Do not adopt it if you want a high-level uncertainty API, since the README itself points to Bayesian Layers and Uncertainty Baselines for that. Before committing, verify that your backend is installed separately, that the tracer stack behaves as you expect under your training loop, and that the experimental directory is not holding code you depend on.

Official sources

1. google/edward2 on GitHub
2. Issues
3. License: Apache-2.0
4. README