[Oryx] Add plate transformation for using named axes in Oryx that parallels

the design of the `Sharded` distribution in TFP.

Also simplifies the `HigherOrderPrimitive`.

PiperOrigin-RevId: 380879501

Author: sharadmv

spinoffs/oryx/oryx/core/ppl/BUILD

@@ -25,6 +25,7 @@ py_library(
     srcs = ["transformations.py"],
     srcs_version = "PY3",
     deps = [
+        ":plate_util",
         # jax dep,
         "//oryx/core:primitive",
         "//oryx/core/interpreters:harvest",
@@ -37,8 +38,19 @@ py_library(
     name = "ppl",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
+    deps = [":transformations"],
+)
+
+# pytype_strict
+py_library(
+    name = "plate_util",
+    srcs = ["plate_util.py"],
+    srcs_version = "PY3",
     deps = [
-        ":transformations",
+        # jax dep,
+        "//oryx/core:primitive",
+        "//oryx/core/interpreters:log_prob",
+        "//oryx/core/interpreters:propagate",
     ],
 )
 
@@ -52,7 +64,9 @@ py_test(
         ":transformations",
         # absl/testing:absltest dep,
         # jax dep,
+        # numpy dep,
         "//oryx/core/interpreters:log_prob",
         "//oryx/internal:test_util",
+        # tensorflow_probability/substrates:jax dep,
     ],
 )

spinoffs/oryx/oryx/core/ppl/__init__.py

@@ -23,6 +23,7 @@
 from oryx.core.ppl.transformations import log_prob
 from oryx.core.ppl.transformations import LogProbFunction
 from oryx.core.ppl.transformations import nest
+from oryx.core.ppl.transformations import plate
 from oryx.core.ppl.transformations import Program
 from oryx.core.ppl.transformations import random_variable
 from oryx.core.ppl.transformations import RANDOM_VARIABLE

spinoffs/oryx/oryx/core/ppl/plate_util.py

@@ -0,0 +1,71 @@
+# Copyright 2021 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Contains utilities for the `plate` transformation.
+
+A plate is a term in graphical models that is used to designate independent
+random variables. In Oryx, `plate` is a transformation that converts a program
+into one that produces independent samples. Ordinarily, this can be done with
+`jax.vmap`, where we could split several random keys and map a program over
+them. Unlike `jax.vmap`, `plate` operates using named axes. A `plate`-ed
+program will specialize the random seed to the particular index of the axis
+being mapped over. Taking the `log_prob` of a `plate` program will reduce over
+the named axis. In design, `plate` resembles the `Sharded` meta-distribution
+from TensorFlow Probability.
+
+In implementation, `plate` is an Oryx `HigherOrderPrimitive` (i.e. a JAX
+`CallPrimitive` with a `log_prob` rule that reduces over a named axis at the
+end.
+"""
+import functools
+
+from jax import lax
+from jax import random
+
+from oryx.core import primitive
+from oryx.core.interpreters import log_prob
+from oryx.core.interpreters import propagate
+
+
+__all__ = [
+    'make_plate',
+]
+
+
+plate_p = primitive.HigherOrderPrimitive('plate')
+
+
+def plate_log_prob_rule(incells, outcells, *, plate, **params):
+  incells, outcells, lp = propagate.call_rule(
+      plate_p, incells, outcells, plate=plate, **params)
+  return incells, outcells, lax.psum(lp, plate)
+
+
+log_prob.log_prob_rules[plate_p] = plate_log_prob_rule
+log_prob.log_prob_registry.add(plate_p)
+
+
+def make_plate(f, *, name):
+  """Wraps a probabilistic program in a plate with a named axis."""
+
+  @functools.wraps(f)
+  def plate_fun(key, *args, **kwargs):
+    key = random.fold_in(key, lax.axis_index(name))
+    return f(key, *args, **kwargs)
+
+  def wrapped(key, *args, **kwargs):
+    return primitive.call_bind(
+        plate_p, plate=name)(plate_fun)(key, *args, **kwargs)
+
+  return wrapped

spinoffs/oryx/oryx/core/ppl/transformations.py

@@ -219,6 +219,7 @@ def f(key):
 from oryx.core import primitive
 from oryx.core.interpreters import harvest
 from oryx.core.interpreters import log_prob as lp
+from oryx.core.ppl import plate_util
 
 __all__ = [
     'block',
@@ -245,7 +246,10 @@ def f(key):
 
 
 @functools.singledispatch
-def random_variable(obj, *, name: Optional[str] = None) -> Program:
+def random_variable(obj,
+                    *,
+                    name: Optional[str] = None,
+                    plate: Optional[str] = None) -> Program:  # pylint: disable=redefined-outer-name
   """A single-dispatch function used to tag values and the outputs of programs.
 
   `random_variable` is a single-dispatch function that enables registering
@@ -255,16 +259,67 @@ def random_variable(obj, *, name: Optional[str] = None) -> Program:
   Args:
     obj: A JAX type to be tagged.
     name (str): A string name to tag input value, cannot be `None`.
+    plate (str): A string named axis for this random variable's plate.
 
   Returns:
     The input value.
   """
   if name is None:
     raise ValueError(f'Cannot call `random_variable` on {type(obj)} '
                      'without passing in a name.')
+  if plate is not None:
+    raise ValueError(f'Cannot call `random_variable` on {type(obj)} '
+                     'with a plate.')
   return harvest.sow(obj, tag=RANDOM_VARIABLE, name=name, mode='strict')
 
 
+def plate(f: Optional[Program] = None, name: Optional[str] = None):
+  """Transforms a program into one that draws samples on a named axis.
+
+  In graphical model parlance, a plate designates independent random variables.
+  The `plate` transformation follows this idea, where a `plate`-ed program
+  draws independent samples. Unlike `jax.vmap`-ing a program, which also
+  produces independent samples with positional batch dimensions, `plate`
+  produces samples with implicit named axes. Named axis support is useful for
+  other JAX transformations like `pmap` and `xmap`.
+
+  Specifically, a `plate`-ed program creates a different key for each axis
+  of the named axis. `log_prob` reduces over the named axis to produce a single
+  value.
+
+  Example usage:
+  ```python
+  @ppl.plate(name='foo')
+  def model(key):
+    return random_variable(random.normal)(key)
+  # We can't call model directly because there are implicit named axes present
+  try:
+    model(random.PRNGKey(0))
+  except NameError:
+    print('No named axis present!')
+  # If we vmap with a named axis, we produce independent samples.
+  vmap(model, axis_name='foo')(random.split(random.PRNGKey(0), 3)) #
+  ```
+
+  Args:
+    f: a `Program` to transform. If `f` is `None`, `plate` returns a decorator.
+    name: a `str` name for the plate which can used as a name axis in JAX
+      functions and transformations.
+
+  Returns:
+    A decorator if `f` is `None` or a transformed program if `f` is provided.
+    The transformed program behaves produces independent across a named
+    axis with name `name`.
+  """
+
+  def transform(f: Program) -> Program:
+    return plate_util.make_plate(f, name=name)
+
+  if f is not None:
+    return transform(f)
+  return transform
+
+
 # Alias for random_variable
 rv = random_variable
 
@@ -273,21 +328,26 @@ def random_variable(obj, *, name: Optional[str] = None) -> Program:
 @random_variable.register(functools.partial)
 def function_random_variable(f: Program,
                              *,
-                             name: Optional[str] = None) -> Program:
+                             name: Optional[str] = None,
+                             plate: Optional[str] = None) -> Program:  # pylint: disable=redefined-outer-name
   """Registers functions with the `random_variable` single dispatch function.
 
   Args:
     f: A probabilistic program.
     name (str): A string name that is used to when tagging the output of `f`.
+    plate (str): A string named axis for this random variable's plate.
 
   Returns:
     A probabilistic program whose output is tagged with `name`.
   """
 
   def wrapped(*args, **kwargs):
+    fun = f
+    if plate is not None:
+      fun = plate_util.make_plate(fun, name=plate)
     if name is not None:
-      return random_variable(nest(f, scope=name)(*args, **kwargs), name=name)
-    return f(*args, **kwargs)
+      return random_variable(nest(fun, scope=name)(*args, **kwargs), name=name)
+    return fun(*args, **kwargs)
 
   return wrapped