[Oryx] Enable vmap rule for tfd.Distributions

Also bumps minor version

PiperOrigin-RevId: 381133393

Author: sharadmv

spinoffs/oryx/oryx/distributions/BUILD

@@ -38,6 +38,7 @@ py_library(
     deps = [
         # jax dep,
         "//oryx/core:primitive",
+        "//oryx/core:trace_util",
         "//oryx/core/interpreters:harvest",
         "//oryx/core/interpreters:log_prob",
         "//oryx/core/interpreters:unzip",

spinoffs/oryx/oryx/distributions/distribution_extensions.py

@@ -14,12 +14,17 @@
 # ============================================================================
 # Lint as: python3
 """Wraps TFP distributions for use with Jax."""
+import itertools as it
+
 from typing import Optional
 
+import jax
 from jax import tree_util
 from jax import util as jax_util
+from jax.interpreters import batching
 from oryx.core import ppl
 from oryx.core import primitive
+from oryx.core import trace_util
 from oryx.core.interpreters import harvest
 from oryx.core.interpreters import inverse
 from oryx.core.interpreters import log_prob
@@ -36,7 +41,7 @@
 
 
 def random_variable_log_prob_rule(flat_incells, flat_outcells, *, num_consts,
-                                  in_tree, out_tree, **_):
+                                  in_tree, out_tree, batch_ndims, **_):
   """Registers Oryx distributions with the log_prob transformation."""
   _, incells = jax_util.split_list(flat_incells, [num_consts])
   val_incells = incells[1:]
@@ -48,13 +53,37 @@ def random_variable_log_prob_rule(flat_incells, flat_outcells, *, num_consts,
   flat_outvals = [cell.val for cell in flat_outcells]
   _, dist = tree_util.tree_unflatten(in_tree, seed_flat_invals)
   outval = tree_util.tree_unflatten(out_tree, flat_outvals)
-  return flat_incells, flat_outcells, dist.log_prob(outval)
+  return flat_incells, flat_outcells, dist.log_prob(outval).sum(
+      axis=list(range(batch_ndims)))
 
 log_prob.log_prob_rules[random_variable_p] = random_variable_log_prob_rule
 
 log_prob.log_prob_registry.add(random_variable_p)
 
 
+def random_variable_batching_rule(args, dims, *, num_consts, batch_ndims,
+                                  jaxpr, **params):
+  """Batching (vmap) rule for the `random_variable` primitive."""
+  old_consts = args[:num_consts]
+  args, dims = args[num_consts:], dims[num_consts:]
+  def _run(*args):
+    return random_variable_p.impl(*it.chain(old_consts, args),
+                                  num_consts=len(old_consts),
+                                  jaxpr=jaxpr,
+                                  batch_ndims=batch_ndims,
+                                  **params)
+  run = jax.vmap(_run, in_axes=dims, out_axes=0)
+  closed_jaxpr, _ = trace_util.stage(run, dynamic=True)(*args)
+  new_jaxpr, new_consts = closed_jaxpr.jaxpr, closed_jaxpr.literals
+  result = random_variable_p.bind(*it.chain(new_consts, args),
+                                  num_consts=len(new_consts),
+                                  jaxpr=new_jaxpr,
+                                  batch_ndims=batch_ndims + 1,
+                                  **params)
+  return result, (0,) * len(result)
+batching.primitive_batchers[random_variable_p] = random_variable_batching_rule
+
+
 def _sample_distribution(key, dist):
   return dist.sample(seed=key)
 
@@ -66,10 +95,15 @@ def distribution_random_variable(dist: tfd.Distribution, *,
   """Converts a distribution into a sampling function."""
   if plate is not None:
     dist = tfed.Sharded(dist, plate)
+  if dist.batch_shape != []:  # pylint: disable=g-explicit-bool-comparison
+    raise ValueError(
+        f'Cannot use a distribution with `batch_shape`: {dist.batch_shape}. '
+        'Instead, use `jax.vmap` or `ppl.plate` to draw independent samples.')
   def wrapped(key):
     def sample(key):
       result = primitive.initial_style_bind(
           random_variable_p,
+          batch_ndims=0,
           distribution_name=dist.__class__.__name__)(_sample_distribution)(
               key, dist)
       return result

spinoffs/oryx/oryx/distributions/distribution_extensions_test.py

@@ -252,6 +252,30 @@ def test_plate_reduces_over_named_axes(self):
     np.testing.assert_allclose(
         tfd.Normal(0., 1.).log_prob(jnp.arange(3.)).sum(), out)
 
+  def test_vmapping_distribution_reduces_to_scalar_log_prob(self):
+
+    def model(key):
+      return jax.vmap(ppl.rv(tfd.Normal(0., 1.)))(random.split(key))
+
+    out = ppl.log_prob(model)(jnp.arange(2.))
+    np.testing.assert_allclose(
+        tfd.Normal(0., 1.).log_prob(jnp.arange(2.)).sum(), out)
+
+  def test_can_map_over_batches_with_vmap_and_reduce_to_scalar_log_prob(self):
+
+    def f(key, x):
+      return ppl.rv(tfd.Normal(x, 1.))(key)
+
+    def model(key, xs):
+      return jax.vmap(f)(random.split(key), xs)
+
+    out = ppl.log_prob(model)(jnp.arange(2.), 2 * jnp.arange(2.))
+    np.testing.assert_allclose(
+        tfd.Normal(jnp.arange(2.), 1.).log_prob(2 * jnp.arange(2.)).sum(), out)
+
+  def test_cannot_use_distribution_with_nontrivial_batch_shape(self):
+    with self.assertRaises(ValueError):
+      ppl.rv(tfd.Normal(jnp.ones(2), 1.))(random.PRNGKey(0))
 
 if __name__ == '__main__':
   absltest.main()

spinoffs/oryx/oryx/version.py

@@ -16,8 +16,8 @@
 
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = '0'
-_MINOR_VERSION = '1'
-_PATCH_VERSION = '4'
+_MINOR_VERSION = '2'
+_PATCH_VERSION = '0'
 
 # When building releases, we can update this value on the release branch to
 # reflect the current release candidate ('rc0', 'rc1') or, finally, the official