v0.10.3

What's Changed

• Kill the manual CHANGELOG (we already have an automated one). by @femtomc in #1564
• [Cold Harbor, 1 / N] Organize the lowest levels of the GenJAX compiler. by @femtomc in #1566
• New logo, docs assets clean up. by @femtomc in #1568
• Update README.md to taste. by @femtomc in #1570
• Bump webfactory/ssh-agent from 0.9.0 to 0.9.1 by @dependabot in #1575
• Add support for sample_shape to primitive distributions by @sritchie in #1576

Full Changelog: v0.10.2...v0.10.3

v0.10.2

What's Changed

• Show mapping tutorial in docs by @fzaiser in #1557
• Do a little cleanup around InitialStylePrimitive. by @femtomc in #1556
• Update usage of JAX partial_eval / wrap_init to accept debug info by @femtomc in #1563

Full Changelog: v0.10.1...v0.10.2

v0.10.1

What's Changed

• Fix cookbook doc by @MathieuHuot in #1539
• Add citation file, research code disclaimer (GEN-1009, GEN-1010) by @sritchie in #1542
• Add CNAME file, fix github pages address by @sritchie in #1545
• [pre-commit.ci] pre-commit autoupdate by @pre-commit-ci in #1544
• Add colab buttons, public colab link to notebook (GEN-1014, GEN-1015) by @sritchie in #1546
• Downgrade poetry (GEN-990) by @sritchie in #1549
• improving cookbook on adding a distribution to GenJAX by @MathieuHuot in #1540
• Fix up new issue templates (GEN-1016) by @sritchie in #1550
• Gen1D cookbook: intro to modelling and block-gibbs by @MathieuHuot in #1541
• add orcids (GEN-1033) by @sritchie in #1547
• Add mapping tutorial to the cookbook (inactive) by @fzaiser in #1552
• Remove EditRequest restriction for DiMap by @fzaiser in #1551
• Update README.md by @femtomc in #1554

New Contributors

• @fzaiser made their first contribution in #1552

Full Changelog: v0.10.0...v0.10.1

v0.10.0

What's Changed

Breaking Changes

• Remove MaskSel, Flag => bool in a number of spots (GEN-979) by @sritchie in #1529
• Expunge the concept of Constraint; just rely on choice map everywhere. by @femtomc in #1532
• Clean up the usage of the trace primitive and remove global "handler stack". by @femtomc in #1531

Fixes

• Fix mix side-effecting issue (GEN-986) by @sritchie in #1537
• use existing index to descend to subtrace in SwitchTrace by @littleredcomputer in #1536

Cookbook

• Adding IndexRequest for Vmap and cookbook entry on IndexRequest by @MathieuHuot in #1518
• Add mathjax support for $, $$ delimiters (GEN-975) by @sritchie in #1525
• fix math rendering in notebooks (GEN-987) by @sritchie in #1535

Misc

• Clarify which keyword arg flip and bernoulli populate in tfd.Bernoulli (GEN-626) by @sritchie in #1526
• warn on implicit use of log scale in categorical, bernoulli by @littleredcomputer in #1528
• Tighten up types on get_submap, get_subtrace (GEN-980) by @sritchie in #1530
• Upgrade to beartype 0.20.0 by @sritchie in #1538
• Remove extra google auth calls by @sritchie in #1522
• Remove docs on deprecated ^ for choicemaps (GEN-961) by @sritchie in #1524

Full Changelog: v0.9.3...v0.10.0

v0.9.3: bugfix release

What's Changed

• Take multiple address components in get_submap by @sritchie in #1516
• Adding a cookbook on speedup using incremental computation by @MathieuHuot in #1499
• equip assess with readable error if address absent by @littleredcomputer in #1517
• Add tfp_distribution back into appropriate init by @sritchie in #1520

Full Changelog: v0.9.2...v0.9.3

v0.9.2: simpler trace printing

What's Changed

• remove info parameter from DimapCombinator by @sritchie in #1511
• remove Combinator suffix from combinators by @sritchie in #1514
• remove ExactDensityFromCallables by @littleredcomputer in #1510
• fix printed repr of genjax distributions created by lambda by @littleredcomputer in #1515

Full Changelog: v0.9.1...v0.9.2

v0.9.1, trace simplification proceeds

What's Changed

• Restrict Selection to static addresses only by @sritchie in #1507
• Remove Projection in favor of Selection by @sritchie in #1509
• GEN-963 move get_subtrace to Trace by @littleredcomputer in #1508

Full Changelog: v0.9.0...v0.9.1

v0.9.0: Trace tidying, we're on pypi!

This is the first GenJAX release published to PyPI proper.

Pypi page: https://pypi.org/project/genjax/
Docs live here: https://chisym.github.io/genjax/

What's Changed

New Stuff

• Adding some TFP distributions (initially Gamma), a test, nuking Bates as it doesn't work by @MathieuHuot in #1493
• Add switch to ChoiceMapBuilder (GEN-897) by @sritchie in #1469
• Add all, none, leaf and S[()] == S.none (GEN-945) by @sritchie in #1492
• Implement project for VmapCombinator by @sritchie in #1505

Changes / Fixes

• Fix Switch choicemap | behavior with static chms (GEN-908) by @sritchie in #1468
• Remove EmptyConstraint and MaskedConstraint as unnecessary concepts by @femtomc in #1477
• Allow kwarg construction of masks (GEN-944) by @sritchie in #1491
• gen-946 deprecate get_sample by @littleredcomputer in #1489
• GEN-957 internal StaticTrace improvements by @littleredcomputer in #1500

Cookbook

• Cleanup cookbook titles after moving away from quarto by @MathieuHuot in #1498
• Adding a cookbook entry on importance method by @MathieuHuot in #1494
• Cookbook on update by @MathieuHuot in #1496
• Move cookbooks out from under quarto, into docs directory. by @femtomc in #1483

Misc

• relax python bounds, upgrade studio by @sritchie in #1497
• Pypi genstudio, publish, bump python to 3.11 by @sritchie in #1476
• Tidy readme for the new pypi world by @sritchie in #1479
• Move docs to GitHub Pages by @sritchie in #1487
• Remove unnecessary gcloud action by @sritchie in #1478
• Fix Codecov integration by @sritchie in #1484
• Replace chi-collective with ChiSym in all docs, badges by @sritchie in #1502
• Fix coverage badge by @sritchie in #1503
• Tweak description for pypi by @sritchie in #1475

Full Changelog: v0.8.1...v0.8.2

v0.8.1

What's Changed

• fix: slack command typo by @cameronfreer in #1457
• Fix non-Diff return value issue (GEN-902) by @sritchie in #1461
• Don't depend on removed jax.interpreters.ad code (GEN-896) by @sritchie in #1465
• Fix assess in vmap (GEN-903) by @sritchie in #1464

New Contributors

• @cameronfreer made their first contribution in #1457

Full Changelog: v0.8.0...v0.8.1

v0.8.0: simpler ChoiceMaps, new combinators

GenJAX 0.8.0 is here! We have a few minor breaking changes, a nice set of new features and a bunch of ChoiceMap improvements for your pleasure. The sections below list out each PR that was merged since the previous release, along with some exposition about the new features where appropriate.

To install this new version, run

pip install keyring keyrings.google-artifactregistry-auth
pip install "genjax==0.8.0" --extra-index-url https://us-west1-python.pkg.dev/probcomp-caliban/probcomp/simple/

or bump the version to "0.8.0" wherever you've pinned your current version.

Thanks to @georgematheos , @esli999 , @MathieuHuot , @femtomc, @limarta and @ahiser1117 for code, testing, bug reports and repros.

Breaking Changes

• dimap's post-xform should accept the original args (GEN-781) by @sritchie in #1429
• Present transformed args to post_fn for DimapCombinator by @sritchie in #1449

dimap's post argument and contramap's argument now accept 3 arguments: the pre-transformed input args, the post-transformed input args, and the return value. Before they only took the pre-transformed args and the return value.

• make genjax.flip return a boolean (GEN-869) by @sritchie in #1451

Fixes / new stuff

• Add exponential distribution. by @femtomc in #1417
• Fix vmap over Pytree (GEN-853) by @sritchie in #1433

This fix allows you to use pytree instances as the input to a vmapped model.

• Add a "static" version of Const (GEN-852) by @sritchie in #1434

This allows you to use Const.unwrap(x) on a Const or non-Const object; Const.unwrap acts as identity when supplied with a non-Const.

• Make @genjax.gen work on methods (GEN-541) by @sritchie in #1349

It's now possible to use the @gen decorator on the method of a class. See below for an example:

In [5]: import jax
   ...: import genjax
   ...: from genjax._src.core.pytree import Pytree
   ...: from genjax._src.core.typing import Array, ArrayLike
   ...:
   ...:
   ...: @genjax.gen
   ...: def outside(x):
   ...:     return genjax.normal(x, 1.0) @ "x"
   ...:
   ...: @Pytree.dataclass
   ...: class Model(Pytree):
   ...:     foo: ArrayLike
   ...:     bar: ArrayLike
   ...:
   ...:     @genjax.gen
   ...:     def run(self, x):
   ...:         y = genjax.normal(self.foo, self.bar) @ "y"
   ...:         z = genjax.normal(x, 1.0) @ "z"
   ...:         return y + z
   ...:
   ...: key = jax.random.PRNGKey(0)
   ...: m = Model(foo=4.0, bar=6.0)
   ...: m.run.simulate(key, (1.0,)).get_choices()
Out[5]:
Static({
  'y': Choice(v=<jax.Array(-4.405305, dtype=float32)>),
  'z': Choice(v=<jax.Array(1.3905154, dtype=float32)>),
})

• masked_iterate, masked_iterate_final (GEN-327, GEN-859) by @sritchie in #1450

We have two new combinators! masked_iterate and masked_iterate_final are similar to their non-masked variants, except instead of providing an initial value, you provide a pair of initial_value and a boolean array. The nth iteration will be scored if the nth boolean entry is True, and not scored if its entry is False.

For example:

masks = jnp.array([True, True])

@genjax.gen
def step(x):
    _ = (
        genjax.normal.mask().vmap(in_axes=(0, None, None))(masks, x, 1.0)
        @ "rats"
    )
    return x

# Create some initial traces:
key = jax.random.key(0)
mask_steps = jnp.arange(10) < 5
model = step.masked_iterate_final()
init_particle = model.simulate(key, (0.0, mask_steps))
init_particle.get_choices()

• (StaticGenerativeFunction) Prevents allocation and accumulation of score, when its available in callee subtraces. by @femtomc in #1418
• Allow PythonicPytree indexing using JAX arrays by @georgematheos in #1440

ChoiceMap

Breaking change — any place you've used an ellipsis like ... to query a ChoiceMap, please replace it with :.

ChoiceMap instances now automatically simplify! There is no more "filtered choicemap", as all filters are pushed down to the leaves:

l_side = C["x"].set({"y":1.0, "z": 2.0}).mask(jnp.array(True))
r_side = C["x", "b"].set(1.0)

l_side | r_side

We've also gotten rid of any returned Indexed choicemap layer. JAX's arrays already keep track of their indexing; there's no need to introduce another Indexed layer on our side.

@genjax.gen
def kernel(mean):
    # Three independent choices
    x = genjax.normal(mean, 1.0) @ "x"
    y = genjax.normal(x, 2.0) @ "y"  # Conditioned on x
    z = genjax.normal(0.0, 3.0) @ "z"
    return x + y + z

vmapped_model = kernel.vmap()
key = jax.random.key(0)
chm = vmapped_model.simulate(key, (jnp.arange(10.0),)).get_choices()
chm

You can still query this choicemap like before, with a : to note the index:

chm[:, "x"]

But you can also just omit the : and get the value via chm["x"]. Any non-string entries you place into the braces will be pushed down to the array-shaped leaves of the choicemap.

• Update choicemap cookbook (GEN-568) by @MathieuHuot in #1381
• Remove ellipsis support from queries, add SwitchChm (GEN-698, GEN-696, GEN-542) by @sritchie in #1421
• Enforce "Mask has scalar flag", allow vectorized Mask (GEN-768, GEN-765, GEN-769) by @sritchie in #1423
• remove Filtered, fix MaskTrace vmappability (GEN-662, GEN-569, GEN-775, GEN-767) by @sritchie in #1422
• Remove Xor, chip away at Or (GEN-774) by @sritchie in #1436
• Remove Indexed layers with None address by @sritchie in #1437

Edit

• (1 / N) Remove ChoiceMapEditRequest, replace with StaticRequest by @femtomc in #1404
• (2 / N) Do a bit of re-org to support deriving the implementations of edit for primitive EditRequest types by @femtomc in #1407
• (3 / N) Add an edit request which can be used to forcibly coerce Diff values according to user input. by @femtomc in #1406
• (5 / N) Add HMC as an EditRequest. by @femtomc in #1408
• (6 / N) Add Rejuvenate -- a proposal-driven regeneration EditRequest by @femtomc in #1411
• (7 / N) Add IndexRequest edit to ScanCombinator by @femtomc in #1405

Cookbook

• Cookbook cleaning splitting by @MathieuHuot in #1380
• clarifying how to use indexed-choicemap by @MathieuHuot in #1425
• convert adev_example notebook to genstudio by @mhuebert in #1414
• update adev_example with latest code from adev_demo by @mhuebert in #1447
• adev: interactive cell (GEN-708) by @mhuebert in #1416

Misc

• Support JAX>=0.4.34 by removing recast_to_float0 and using our own implementation. by @femtomc in #1410
• Adding a test for mask scan combinator by @MathieuHuot in #1443

Full Changelog: v0.7.0...v0.8.0