[Oryx] Update propagate transformation to support state at each equation

PiperOrigin-RevId: 380879404

Author: sharadmv

spinoffs/oryx/oryx/core/interpreters/BUILD

@@ -53,6 +53,7 @@ py_library(
     srcs = ["propagate.py"],
     srcs_version = "PY3",
     deps = [
+        ":harvest",
         # dataclasses dep,
         # jax dep,
         "//oryx/core:pytree",

spinoffs/oryx/oryx/core/interpreters/inverse/core.py

@@ -18,17 +18,13 @@
 import jax
 from jax import abstract_arrays
 from jax import core as jax_core
-from jax import linear_util as lu
 from jax import tree_util
 from jax import util as jax_util
-from jax.interpreters import partial_eval as pe
 from jax.interpreters import pxla
-from jax.interpreters import xla
 import jax.numpy as np
 
 from oryx.core import primitive
 from oryx.core import trace_util
-from oryx.core.interpreters import harvest
 from oryx.core.interpreters import propagate
 from oryx.core.interpreters.inverse import slice as slc
 
@@ -177,8 +173,8 @@ def wrapped(*args, **kwargs):
         for arg in flat_forward_args]
     flat_incells = [InverseAndILDJ.unknown(aval) for aval in flat_forward_avals]
     flat_outcells = safe_map(InverseAndILDJ.new, flat_args)
-    env = propagate.propagate(InverseAndILDJ, ildj_registry, jaxpr.jaxpr,
-                              flat_constcells, flat_incells, flat_outcells)  # pytype: disable=wrong-arg-types
+    env, _ = propagate.propagate(InverseAndILDJ, ildj_registry, jaxpr.jaxpr,
+                                 flat_constcells, flat_incells, flat_outcells)  # pytype: disable=wrong-arg-types
     flat_incells = [env.read(invar) for invar in jaxpr.jaxpr.invars]
     if any(not flat_incell.top() for flat_incell in flat_incells):
       raise ValueError('Cannot invert function.')
@@ -246,6 +242,9 @@ def __getitem__(self, prim):
   def __setitem__(self, prim, val):
     self.rules[prim] = val
 
+  def __contains__(self, prim):
+    return prim in self.rules
+
 
 def register_elementwise(prim):
   """Registers an elementwise primitive with ILDJ."""
@@ -296,46 +295,19 @@ def ildj_rule(incells, outcells, **params):
 ildj_registry = InverseDict()
 
 
-@lu.transformation_with_aux
-def flat_propagate(tree, *flat_invals):
-  invals, outvals = tree_util.tree_unflatten(tree, flat_invals)
-  subenv = yield ((invals, outvals), {})
-  subenv_vals, subenv_tree = tree_util.tree_flatten(subenv)
-  yield subenv_vals, subenv_tree
-
-
-def call_ildj(prim, incells, outcells, **params):
-  """InverseAndILDJ rule for call primitives."""
-  f, incells = incells[0], incells[1:]
-  flat_vals, in_tree = tree_util.tree_flatten((incells, outcells))
-  new_params = dict(params)
-  if 'donated_invars' in params:
-    new_params['donated_invars'] = (False,) * len(flat_vals)
-  f, aux = flat_propagate(f, in_tree)
-  subenv_vals = prim.bind(f, *flat_vals, **new_params)
-  subenv_tree = aux()
-  subenv = tree_util.tree_unflatten(subenv_tree, subenv_vals)
-  new_incells = [subenv.read(var) for var in subenv.jaxpr.invars]
-  new_outcells = [subenv.read(var) for var in subenv.jaxpr.outvars]
-  return new_incells, new_outcells, subenv
-ildj_registry[xla.xla_call_p] = jax_util.partial(call_ildj, xla.xla_call_p)
-ildj_registry[jax_core.call_p] = jax_util.partial(call_ildj, jax_core.call_p)
-ildj_registry[pe.remat_call_p] = jax_util.partial(call_ildj, pe.remat_call_p)
-ildj_registry[harvest.nest_p] = jax_util.partial(call_ildj, harvest.nest_p)
-
-
 def hop_inverse_rule(prim):
-  ildj_registry[prim] = jax_util.partial(call_ildj, prim)
+  ildj_registry[prim] = jax_util.partial(propagate.call_rule, prim)
 primitive.register_hop_transformation_rule('inverse', hop_inverse_rule)
 
 
 def initial_ildj(incells, outcells, *, jaxpr, num_consts, **_):
   const_cells, incells = jax_util.split_list(incells, [num_consts])
-  env = propagate.propagate(InverseAndILDJ, ildj_registry, jaxpr, const_cells,
-                            incells, outcells)  # pytype: disable=wrong-arg-types
+  env, state = propagate.propagate(
+      InverseAndILDJ, ildj_registry, jaxpr, const_cells,
+      incells, outcells)  # pytype: disable=wrong-arg-types
   new_incells = [env.read(invar) for invar in jaxpr.invars]
   new_outcells = [env.read(outvar) for outvar in jaxpr.outvars]
-  return const_cells + new_incells, new_outcells, None
+  return const_cells + new_incells, new_outcells, state
 
 
 def initial_inverse_rule(prim):
@@ -371,7 +343,7 @@ def remove_slice(cell):
   mapped_incells = safe_map(remove_slice, incells)
   mapped_outcells = safe_map(remove_slice, outcells)
   flat_vals, in_tree = tree_util.tree_flatten((mapped_incells, mapped_outcells))
-  f, aux = flat_propagate(f, in_tree)
+  f, aux = propagate.flat_propagate(f, in_tree)
   # Assume all invars as mapped
   new_in_axes = (0,) * len(flat_vals)
   new_params = dict(params, in_axes=new_in_axes)
@@ -383,14 +355,13 @@ def remove_slice(cell):
         lambda: (0,) * aux().num_leaves,
         closure=('ildj', params['out_axes']))
     del new_params['out_axes']
-  subenv_vals = prim.bind(f, *flat_vals, **new_params)
-  subenv_tree = aux()
-  subenv = tree_util.tree_unflatten(subenv_tree, subenv_vals)
-  new_incells = [subenv.read(var) for var in subenv.jaxpr.invars]
-  new_outcells = [subenv.read(var) for var in subenv.jaxpr.outvars]
+  flat_out = prim.bind(f, *flat_vals, **new_params)
+  out_tree = aux()
+  new_incells, new_outcells, state = tree_util.tree_unflatten(
+      out_tree, flat_out)
   new_incells = [add_slice(v, old_v)
                  for old_v, v in safe_zip(incells, new_incells)]
   new_outcells = [add_slice(v, old_v)
                   for old_v, v in safe_zip(outcells, new_outcells)]
-  return new_incells, new_outcells, subenv
+  return new_incells, new_outcells, state
 ildj_registry[pxla.xla_pmap_p] = jax_util.partial(map_ildj, pxla.xla_pmap_p)

spinoffs/oryx/oryx/core/interpreters/log_prob.py

@@ -17,7 +17,6 @@
 from jax import core as jax_core
 from jax import random
 from jax import tree_util
-import jax.numpy as np
 
 from oryx.core import trace_util
 from oryx.core.interpreters import inverse
@@ -47,16 +46,17 @@ def __missing__(self, prim):
     self[prim] = rule = make_default_rule(prim)
     return rule
 
-log_prob_rules = LogProbRules()
 
+log_prob_rules = LogProbRules()
 
 # The log_prob_registry is used to compute log_prob values from samples after
 # propagation is done.
-log_prob_registry = {}
+log_prob_registry = set()
 
 
 def log_prob(f):
   """LogProb function transformation."""
+
   def wrapped(sample, *args, **kwargs):
     """Function wrapper that takes in log_prob arguments."""
     # Trace the function using a random seed
@@ -69,50 +69,55 @@ def wrapped(sample, *args, **kwargs):
         InverseAndILDJ.unknown(trace_util.get_shaped_aval(dummy_seed))
     ] + [InverseAndILDJ.new(val) for val in flat_inargs]
     flat_outcells = [InverseAndILDJ.new(a) for a in flat_outargs]
-    # Re-use the InverseAndILDJ propagation but silently fail instead of
-    # erroring when we hit a primitive we can't invert.
-    env = propagate.propagate(InverseAndILDJ, log_prob_rules, jaxpr.jaxpr,
-                              constcells, flat_incells, flat_outcells)
-    # Traverse the resulting environment, looking for primitives that have
-    # registered log_probs.
-    final_log_prob = _accumulate_log_probs(env)
-    return final_log_prob
-  return wrapped
+    return log_prob_jaxpr(jaxpr.jaxpr, constcells, flat_incells, flat_outcells)
 
+  return wrapped
 
-def _accumulate_log_probs(env):
-  """Recursively traverses Jaxprs to accumulate log_prob values."""
-  final_log_prob = 0.0
-  eqns = safe_map(propagate.Equation.from_jaxpr_eqn, env.jaxpr.eqns)
-  for eqn in eqns:
-    if eqn.primitive in log_prob_registry:
-      var, = eqn.outvars
-      if var not in env:
-        raise ValueError('Cannot compute log_prob of function.')
-      incells = [env.read(v) for v in eqn.invars]
-      outcells = [env.read(v) for v in eqn.outvars]
-      outcell, = outcells
-      if not outcell.top():
-        raise ValueError('Cannot compute log_prob of function.')
-      lp = log_prob_registry[eqn.primitive](
-          [cell if not cell.top() else cell.val for cell in incells],
-          outcell.val, **eqn.params
-      )
-      assert np.ndim(lp) == 0, 'log_prob must return a scalar.'
-      # Accumulate ILDJ term
-      final_log_prob += lp + np.sum(outcell.ildj)
-  for subenv in env.subenvs.values():
-    sub_lp = _accumulate_log_probs(subenv)
-    final_log_prob += sub_lp
+failed_log_prob = object()  # sentinel for being unable to compute a log_prob
+
+
+def log_prob_jaxpr(jaxpr, constcells, flat_incells, flat_outcells):
+  """Runs log_prob propagation on a Jaxpr."""
+
+  def reducer(env, eqn, curr_log_prob, new_log_prob):
+    if curr_log_prob is failed_log_prob or new_log_prob is failed_log_prob:
+      # If `curr_log_prob` is `None` that means we were unable to compute
+      # a log_prob elsewhere, so the propagate failed.
+      return failed_log_prob
+    if eqn.primitive in log_prob_registry and new_log_prob is None:
+      # We are unable to compute a log_prob for this primitive.
+      return failed_log_prob
+    if new_log_prob is not None:
+      cells = [env.read(var) for var in eqn.outvars]
+      ildjs = sum([cell.ildj.sum() for cell in cells if cell.top()])
+      return curr_log_prob + new_log_prob + ildjs
+    return curr_log_prob
+
+  # Re-use the InverseAndILDJ propagation but silently fail instead of
+  # erroring when we hit a primitive we can't invert. We accumulate the log
+  # probability values using the propagater state.
+  _, final_log_prob = propagate.propagate(
+      InverseAndILDJ,
+      log_prob_rules,
+      jaxpr,
+      constcells,
+      flat_incells,
+      flat_outcells,
+      reducer=reducer,
+      initial_state=0.)
+  if final_log_prob is failed_log_prob:
+    raise ValueError('Cannot compute log_prob of function.')
   return final_log_prob
 
 
 def make_default_rule(prim):
   """Creates rule for prim without a registered log_prob."""
+
   def rule(incells, outcells, **params):
     """Executes the inverse rule but fails if the inverse isn't implemented."""
     try:
       return ildj_registry[prim](incells, outcells, **params)
     except NotImplementedError:
       return incells, outcells, None
+
   return rule

spinoffs/oryx/oryx/core/interpreters/log_prob_test.py

@@ -28,6 +28,7 @@
 from oryx.core import state
 from oryx.core.interpreters.log_prob import log_prob
 from oryx.core.interpreters.log_prob import log_prob_registry
+from oryx.core.interpreters.log_prob import log_prob_rules
 from oryx.internal import test_util
 
 random_normal_p = jax_core.Primitive('random_normal')
@@ -49,10 +50,14 @@ def random_normal_abstract(_, name=None):
 random_normal_p.def_abstract_eval(random_normal_abstract)
 
 
-def random_normal_log_prob(_, outval, name=None):
-  del name
-  return bd.Normal(0., 1.).log_prob(outval)
-log_prob_registry[random_normal_p] = random_normal_log_prob
+def random_normal_log_prob_rule(incells, outcells, **_):
+  outcell, = outcells
+  if not outcell.top():
+    return incells, outcells, None
+  outval = outcell.val
+  return incells, outcells, bd.Normal(0., 1.).log_prob(outval)
+log_prob_rules[random_normal_p] = random_normal_log_prob_rule
+log_prob_registry.add(random_normal_p)
 
 
 def call(f):