Edward2: a low-level probabilistic programming language for TensorFlow, JAX and NumPy

What Edward2 actually provides, and who the README is talking to

Edward2 describes itself as a simple probabilistic programming language, and the word simple is doing real work here. What ships in the edward2/ directory is a set of core utilities: random variables, a tracer that intercepts function calls, and a collection of layers. There is no inference engine in the box. The README states that the core utilities are fairly low-level and directs anyone who wants a high-level module for uncertainty modeling to the Bayesian Layers guide, and anyone who wants research-ready code to Uncertainty Baselines. That is an unusually direct way of telling you the library is not the thing you build a product on.

The audience follows from that. Edward2 is for people who are already fluent in TensorFlow Probability distributions or in Flax, and who want to declare a model as a Python function that samples, then transform that function for training or inference. If you want to specify a Bayesian model in a few lines and get posterior samples back, this is the wrong entry point. If you want to write the sampling and the inference yourself, with primitives that compose with the rest of the deep learning stack, it is the right one.

The repository layout reinforces the split. Library code sits in edward2/, worked examples in examples/ (including no_u_turn_sampler/ and notebooks/), and experimental/ is described in the README as active research projects. Code under experimental/ should be read as research output, not as a supported API surface.

Random variables that behave like tensors

The central object is the RandomVariable. According to the README, a random variable carries a probability distribution accessible as rv.distribution, which is a TensorFlow Distribution instance governing methods such as log_prob and sample. Random variables are constructed the way TensorFlow Distributions are constructed, by naming a distribution and passing its parameters.

The design decision worth noticing is that instantiating a random variable creates a sampling op by default. The README says the default number of samples, controllable through the sample_shape argument, is one, and that if you pass the optional value argument no sampling op is created. That last detail matters more than it looks: supplying value is how you pin a random variable to observed data instead of drawing from the prior. There is no separate observed=True flag to remember.

Interoperation is the other half. TensorFlow ops applied to a random variable operate on its sample, so x + y, x / y and tf.tanh(x * y) all return tensors rather than random variables, while indexing such as x[1] returns a random variable again. If you have used a probabilistic language where random variables are opaque objects that must be unwrapped before arithmetic, this is a different feel. It also means the boundary between a random variable and a tensor is easy to cross by accident, and the README's examples are the clearest documentation of where that boundary sits.

Models as functions, and a second function for the posterior

A probabilistic model in Edward2 is a plain Python function. The README's example is Bayesian logistic regression: a function named logistic_regression takes features, draws coeffs and intercept from ed.Normal, feeds them through ed.Bernoulli with logits computed by tensordot plus the intercept, and returns the outcomes random variable. Executing the function runs the generative process. Inputs to the function are what the model conditions on.

The more interesting pattern is the second function in the README, logistic_regression_posterior. It does not model data at all. It constructs a learnable distribution intended to approximate the logistic regression posterior, using ed.MultivariateNormalTriL for the coefficients with a scale_tril built by tfp.trainable_distributions.tril_with_diag_softplus_and_shift, and ed.Normal for the intercept with a softplus-transformed scale plus 1e-5. The parameters are tf.Variable objects created outside the function.

That pair of functions is the whole workflow in miniature. One function says how data could have been generated. The other says what shape your approximation takes. Edward2 gives you the vocabulary for both and stops there; the loss, the optimizer and the training loop are yours to write. If you were hoping the library would connect those two functions for you, the README does not claim it will.

Tracing: rewriting a program's computation

Tracing is the mechanism that makes the rest possible. The README describes a tracer as a function that acts on another function f and its arguments, performs computations, and returns an output, typically f(*args, **kwargs). The ed.trace context manager pushes tracers onto a stack, and any traceable function is intercepted by that stack. The README states that all random variable constructors are traceable.

This is the part of Edward2 that is genuinely a language feature rather than a convenience wrapper. Because random variable construction goes through the stack, a tracer can observe or replace each draw without the model function knowing. That is how the same generative function can serve sampling, scoring, or a transformed computation, depending on which tracer is active. It is also why the library can stay small: the flexibility lives in the interception mechanism instead of in a large set of model classes.

The cost is conceptual. Tracers are a control-flow idea, and debugging a stack of them is not the same as reading a linear program. The README introduces the concept but does not enumerate the tracers that ship in the repository, so expect to read edward2/tracer.py and the examples rather than rely on the prose. Anyone who has only used declarative probabilistic DSLs will find this closer to writing an interpreter than to writing a model.

Installing Edward2 and running a first program

The README recommends the latest development version over the stable release, and gives the reason: the stable version is very rarely updated, described as a passion project maintained by part-timers. Take that at face value. Installing from git is the documented default path.

pip install "edward2 @ git+https://github.com/google/edward2.git"

Installing the package does not install a backend. The README is explicit that Edward2 supports TensorFlow (the default), JAX and NumPy, and that you add the dependency yourself through an extra, for example edward2[tensorflow], replacing tensorflow for the appropriate backend.

pip install edward2[tensorflow]

The setup.py extras confirm the pins: the tensorflow extra requires tensorflow>=2.0.0a0 and tensorflow-probability>=0.8.0, the jax extra requires jax>=0.2.13 and flax>=0.3.4, and the numpy extra requires numpy>=1.7 and scipy>=1.0.0. There is also a tf-nightly extra for tf-nightly and tfp-nightly, which the README says you need when Edward2 uses the latest changes from TensorFlow. Only tqdm is an unconditional dependency.

With a backend in place, the smallest useful program is a random variable and a log probability, exactly as the README shows:

import edward2 as ed

normal_rv = ed.Normal(loc=0., scale=1.)
normal_rv.distribution.log_prob(1.231)

The first line draws one sample from a standard normal. The second evaluates the log density at 1.231 and returns a TensorFlow tensor, not a random variable. If that runs, the backend and the distribution library are wired correctly. The README's example output shows a float32 scalar; your value will differ because it is a draw.

Where Edward2 is the wrong tool

The clearest limitation is the one the README states itself: there is no high-level uncertainty modeling module. Bayesian Layers is a separate guide, and Uncertainty Baselines is recommended for research-ready code. If your task is to add uncertainty to an existing network, starting at Edward2 means assembling the pieces yourself.

Release cadence is the second constraint. The README says the stable version is very rarely updated and that scheduling releases sucks up time for a part-time team. The package's own classifier in setup.py marks it Development Status :: 4 - Beta. Neither of those is disqualifying, but together they mean you should not expect a stable release line to track TensorFlow or JAX changes. The tf-nightly extra exists precisely because Edward2 sometimes depends on unreleased TensorFlow behaviour, which is a coupling that can break your environment when either side moves.

The third is scope. Edward2 gives you random variables and tracing. It does not give you a sampler, a variational objective, or a training loop. The examples/ directory has a no_u_turn_sampler/ example and notebooks, which suggests the intended path is to read and adapt code rather than call a function. If your team needs a documented inference API with a support contract, this is not it.

Edward2 against Pyro and NumPyro

The natural comparison is Pyro or NumPyro, which are also probabilistic programming languages embedded in Python, but the split is architectural rather than cosmetic. Pyro and NumPyro ship an inference layer: effect handlers for sampling statements, and a set of variational and MCMC algorithms you can point at a model. Edward2 ships the random variable and the tracer, and leaves inference to you or to a companion project.

The backend story differs too. Edward2 supports TensorFlow, JAX and NumPy through extras, so a model can move between ecosystems without a rewrite, at least in principle. NumPyro is JAX-only by design and leans on that single backend for its compilation story. If your stack is already JAX and you want inference included, NumPyro is the shorter path. If your stack is TensorFlow and you want random variables that participate in TensorFlow graphs, Edward2 fits where a JAX-only library does not.

The honest summary is that Edward2 is a lower layer than either. Choosing it means choosing to write the inference code. That is a reasonable choice for research that needs to manipulate a model's computation in ways a fixed inference API would not allow, and a poor choice for anyone who wants a posterior without writing the machinery.

Editorial conclusion

Adopt Edward2 if you already know TensorFlow Probability or Flax and want random variables that behave like tensors inside an existing model, or if you are reproducing a paper that builds on Edward2's tracing and random variable primitives. Do not adopt it expecting a sampler, a variational inference loop or a Keras-style fit method; the README points at Bayesian Layers and Uncertainty Baselines for that layer of work. Before committing, check that the backend extra you need installs cleanly against your TensorFlow or JAX version, and read the Upgrading_From_Edward_To_Edward2.md guide if you are porting Edward 1 code, because the two APIs are not interchangeable.

Frequently asked questions

Official sources

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