Skip to content

v1.0.14

Latest

Choose a tag to compare

@femtomc femtomc released this 27 Sep 04:40

GenJAX 1.0.14 is the first 1.x release on PyPI. It replaces 0.10.3, the last release of the 0.x codebase, as the version pip install genjax selects.

Install

  • pip install genjax installs the package.
  • pip install "genjax[viz]" also installs matplotlib for the plotting helpers.
  • The 1.x API differs from 0.x. A project written against 0.x can pin genjax<1.
  • 1.0.14 requires Python 3.12 or later. On Python 3.11, pip install genjax still selects 0.10.3.
  • genjax 0.x depended on tensorflow-probability, and 1.0.14 depends on tfp-nightly. Both install the tensorflow_probability module, so uninstall tensorflow-probability before upgrading an environment from 0.x. To repair an environment upgraded without that step, remove tensorflow-probability and then run pip install --force-reinstall --no-deps tfp-nightly==0.26.0.dev20260831.

Changes since 1.0.13

  • The repository moved from github.com/femtomc/genjax to github.com/a-tiny-project/genjax. github.com/femtomc/genjax and github.com/probcomp/genjax are now forks of it.
  • import genjax no longer needs matplotlib. The plotting names (genjax.viz, raincloud, horizontal_raincloud) load on first use.
  • from genjax import * raised AttributeError in 1.0.13, because __all__ named three functions the package never bound. genjax now exports the MCMC diagnostics compute_rhat and compute_ess, and __all__ no longer lists linear_gaussian_ssm, which existed nowhere.
  • The PyPI page shows the description, the Apache-2.0 license, keywords, and links to the repository, the paper DOI, and the Zenodo DOI. The README's links resolve on PyPI.
  • The source distribution holds the package, its tests, the examples, and the citation file.
  • Two case studies are ported from the artifact of the PLDI 2024 paper "Probabilistic Programming with Programmable Variational Inference": the cone model's variational objectives (its Table 4) and the AIR estimators (attend-infer-repeat, a model that explains an image as a variable number of objects).
  • In ADEV (automatic differentiation of expected values, which supplies the gradient estimators for variational inference), the flip primitives sample through the flip distribution, so every estimator reads its argument as a probability. ADEV treats JAX float0 tangents (the tangent type of integer and boolean values) as zero. genjax exports two new ADEV estimators for uniform draws: uniform_reinforce (score function) and uniform_reparam (reparameterization).
  • init_csmc and extend_csmc give particle 0 the retained latent choices, the observed values, and the same importance weight as the other particles. In 1.0.13 init_csmc kept a proposal's trace for particle 0 and weighted it by the joint density, which biased conditional SMC, and extend_csmc raised TypeError on every call.
  • With the default proposal, conditional SMC requires that no latent choice's distribution reads an observed value. A model where one does needs a custom proposal.
  • seed, modular_vmap and the other transforms that stage a function accept functions with return annotations. In 1.0.13, seed raised TypeError on a function annotated to return an array.
  • An exception raised inside a generative function no longer leaves its handler installed. In 1.0.13, after such an exception, later distribution calls and ADEV sites in the same process, such as the next notebook cell, could return Thunk objects instead of values.
  • The requirements changed:
    • jax>=0.11.1,<0.12 is now declared. 1.0.13 declared no JAX requirement.
    • TensorFlow Probability comes from tfp-nightly==0.26.0.dev20260831 instead of tensorflow-probability>=0.25.0,<0.26.
    • beartype moves from >=0.21.0,<0.22 to >=0.22.9,<0.23.
    • jaxtyping moves from >=0.3.2,<0.4 to >=0.3.11,<0.4.
  • Deprecated command-line paths and compatibility shims are removed.
  • Comments and docstrings were trimmed. Those edits change no code.
  • In the Game of Life example, create_showcase_figure(load_from_file=...) returns its figure. In 1.0.13 it raised UnboundLocalError on every load.
  • The Game of Life showcase defaults to the wizards pattern, and the mit, popl and hermes patterns are removed. In 1.0.13 the default mit pattern raised NameError, and the images for those three patterns were not in the repository. The wizards image is no longer in the repository or the source distribution either. The README gives its pinned URL and SHA-256.
  • In the localization example, plot_smc_method_comparison(include_legend=True) draws its legend. In 1.0.13 it raised NameError.
  • Development, the examples and the performance benchmarks use uv dependency groups and a committed uv.lock for Linux x86_64 and macOS arm64. 1.0.13 used Pixi environments. The move to uv changes no requirement of the published package.
  • The performance benchmark's importance sampling workload now matches across the compared frameworks.

Repository history

The commit messages on main were edited on 2026-09-25, and on 2026-09-27 the tags v0.1.0 through v1.0.13 moved to the edited commits. Every commit and every tag kept its file tree. A clone made before those dates needs git fetch origin --tags --force and then git reset --hard origin/main, which discards any local commits on main.

Acknowledgments

GenJAX 1.0 continues the GenJAX project, whose 0.x releases were developed in genjax-community/genjax from 2022 to 2025. GenJAX thanks the 22 people other than the maintainer who contributed commits to that codebase:

Matthew Brulhardt, Jacob Burnim, Guillaume Dalle, Arijit Dasgupta, Cameron Freer, Matin Ghavami, Alex Hiser, Matt Huebert, Mathieu Huot, Mirko Klukas, Urs Köster, Ben Lee, Ian Limarta, Joao Loula, David R. MacIver, George Matheos, Jay Pottharst, Sam Ritchie, Rif A. Saurous, Colin Smith, Xiaoyan Wang, Fabian Zaiser.

Files

The wheel and the source distribution are attached. PyPI serves the same files. Their SHA-256 checksums:

04634bce70015dae58586afa06750e6b095b89f162efb6a15f7de1f1eda7ed3b  genjax-1.0.14-py3-none-any.whl
e6a73a2994bff442461d9f6c3db7c7aa0677f22f14864e94c607bfaa0c01e696  genjax-1.0.14.tar.gz

Zenodo archives this release under the GenJAX record, and CITATION.cff gives the citation.