Simplify sharded JDs by removing reduce_over_shards

PiperOrigin-RevId: 375595276

Author: sharadmv

tensorflow_probability/python/experimental/distribute/__init__.py

@@ -21,13 +21,16 @@
 from tensorflow_probability.python.experimental.distribute.joint_distribution import JointDistributionNamed
 from tensorflow_probability.python.experimental.distribute.joint_distribution import JointDistributionSequential
 from tensorflow_probability.python.experimental.distribute.sharded import Sharded
+from tensorflow_probability.python.internal.distribute_lib import make_pbroadcast_function
 from tensorflow_probability.python.internal.distribute_lib import make_psum_function
 from tensorflow_probability.python.internal.distribute_lib import make_sharded_log_prob_parts
 
 __all__ = [
     'JointDistributionCoroutine',
     'JointDistributionNamed',
     'JointDistributionSequential',
+    'make_pbroadcast_function',
+    'make_psum_function',
     'make_sharded_log_prob_parts',
     'Sharded',
 ]

tensorflow_probability/python/experimental/distribute/joint_distribution.py

@@ -18,8 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import functools
-
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python import distributions as distribution_lib
 from tensorflow_probability.python.distributions import log_prob_ratio as lp_ratio
@@ -39,30 +37,23 @@ def _map_measure_over_dists(self, attr, value):
     if (attr in ('log_prob', 'unnormalized_log_prob')) and any(
         self.experimental_shard_axis_names):
 
-      def inner_log_prob_parts(flat_value):
-        unflat_value = self._model_unflatten(flat_value)
+      def inner_log_prob_parts(value):
         ds, xs = self._call_flat_sample_distributions(
-            value=unflat_value, seed=samplers.zeros_seed())
-        # For sharded distributions, we need to make sure not to do an
-        # all-reduce.
-        axis_names = self._model_flatten(self.experimental_shard_axis_names)
-        log_prob_fns = [
-            functools.partial(getattr(d, attr), reduce_over_shards=False)
-            if axis_name else getattr(d, attr)
-            for d, axis_name in zip(ds, axis_names)
-        ]
+            value=value, seed=samplers.zeros_seed())
         # We need to flatten and unflatten here to ensure the output structure
         # matches `flat_sharded_distributions`.
-        vals = self._model_unflatten(
-            [log_prob_fn(x) for log_prob_fn, x in zip(log_prob_fns, xs)])
-        return self._model_flatten(vals)
-
-      flat_value = self._model_flatten(value)
-      flat_axis_names = self._model_flatten(self.experimental_shard_axis_names)
-      flat_xs = distribute_lib.make_sharded_log_prob_parts(
-          inner_log_prob_parts, flat_axis_names)(
-              flat_value)
-      return iter(flat_xs)
+        return self._model_unflatten(
+            [getattr(d, attr)(x) for d, x in zip(ds, xs)])
+
+      axis_names = self.experimental_shard_axis_names
+      # Individual distributions will apply psum in their `log_prob` methods
+      # so we need to pbroadcast `value` according to `axis_names` to provide
+      # correct gradients. We are safe to add pbroadcasts to functions with
+      # psums already in them.
+      log_prob_parts = distribute_lib.make_pbroadcast_function(
+          inner_log_prob_parts, (axis_names,), axis_names,
+          out_dtype=value)(value)
+      return iter(tf.nest.flatten(log_prob_parts))
     ds, xs = self._call_flat_sample_distributions(
         value=value, seed=samplers.zeros_seed())
     return (getattr(d, attr)(x) for d, x in zip(ds, xs))
@@ -104,16 +95,14 @@ def _dist_jd_log_prob_ratio(p, x, q, y, name=None):
     def log_prob_ratio_parts_fn(x, y):
       p_dists = p.sample_distributions(value=x, seed=samplers.zeros_seed())[0]
       q_dists = q.sample_distributions(value=y, seed=samplers.zeros_seed())[0]
-      # Ensure sharded distributions defer reductions.
-      kwds = lambda a: {'reduce_over_shards': False} if a else {}
       return nest.map_structure_up_to(
           p_dists,
-          lambda p, x, q, y, s: lp_ratio.log_prob_ratio(p, x, q, y, **kwds(s)),
-          p_dists, x, q_dists, y, p_axis_names)
+          lp_ratio.log_prob_ratio,
+          p_dists, x, q_dists, y)
 
     return tf.add_n(
         tf.nest.flatten(
-            distribute_lib.make_psum_function(
+            distribute_lib.make_pbroadcast_function(
                 log_prob_ratio_parts_fn,
                 in_axes=(p_axis_names, p_axis_names),
                 out_axes=p_axis_names,

tensorflow_probability/python/experimental/distribute/joint_distribution_test.py

@@ -155,6 +155,7 @@ def run(key):
 
     keys = tfp.random.split_seed(self.key, 2)
     samples = []
+    unmapped_samples = []
     log_probs = []
     true_log_probs = []
 
@@ -172,24 +173,51 @@ def run(key):
       true_log_prob = true_log_prob_fn(w, x, data)
 
       samples.append(sample)
+      unmapped_samples.append((w, x, data))
       log_probs.append(log_prob[0])
       true_log_probs.append(true_log_prob)
 
+    def true_diff(x, y):
+      return true_log_prob_fn(*x) - true_log_prob_fn(*y)
+
     def run_diff(x, y):
-      return tfp.experimental.distributions.log_prob_ratio(dist, x, dist, y)
+      def _lpr(x, y):
+        return tfp.experimental.distributions.log_prob_ratio(dist, x, dist, y)
+      return tfp.math.value_and_gradient(_lpr, [x, y])
 
-    dist_lp_diff = self.per_replica_to_tensor(
+    dist_lp_diff, dist_lp_diff_grad = self.per_replica_to_tensor(
         self.strategy_run(
             run_diff, tuple(tf.nest.map_structure(self.shard_values, samples))))
 
-    true_lp_diff = true_log_probs[0] - true_log_probs[1]
+    true_lp_diff, true_lp_diff_grad = tfp.math.value_and_gradient(
+        true_diff, unmapped_samples)
+
+    if isinstance(dist, jd.JointDistributionNamed):
+      dist_lp_diff_grad[0] = (
+          dist_lp_diff_grad[0]['w'][0],
+          dist_lp_diff_grad[0]['x'],
+          dist_lp_diff_grad[0]['data'])
+      dist_lp_diff_grad[1] = (
+          dist_lp_diff_grad[1]['w'][0],
+          dist_lp_diff_grad[1]['x'],
+          dist_lp_diff_grad[1]['data'])
+    else:
+      true_lp_diff_grad[0] = list(true_lp_diff_grad[0])
+      true_lp_diff_grad[1] = list(true_lp_diff_grad[1])
+      dist_lp_diff_grad[0] = list(dist_lp_diff_grad[0])
+      dist_lp_diff_grad[0][0] = dist_lp_diff_grad[0][0][0]
+      dist_lp_diff_grad[1] = list(dist_lp_diff_grad[1])
+      dist_lp_diff_grad[1][0] = dist_lp_diff_grad[1][0][0]
+
     lp_diff = log_probs[0] - log_probs[1]
 
     self.assertAllClose(
-        self.evaluate(true_lp_diff), self.evaluate(lp_diff),
+        true_lp_diff, lp_diff,
         rtol=7e-6)  # relaxed tol for fp32 in JAX
     self.assertAllClose(
-        self.evaluate(true_lp_diff), self.evaluate(dist_lp_diff[0]))
+        true_lp_diff, dist_lp_diff[0])
+    self.assertAllClose(
+        true_lp_diff_grad, dist_lp_diff_grad)
 
   def test_default_event_space_bijector_non_interacting(self):
 

tensorflow_probability/python/experimental/distribute/sharded.py

@@ -35,16 +35,9 @@
 
 def _implement_sharded_lp_fn(fn_name):
   """Implements log_prob or unnormalized_log_prob."""
-  def lp_fn(self, x, reduce_over_shards=True, **kwargs):
-
-    new_kwargs = dict(kwargs)
-    if self.distribution.experimental_shard_axis_names:
-      new_kwargs['reduce_over_shards'] = reduce_over_shards
-    lp = getattr(self.distribution, fn_name)(x, **new_kwargs)
-    if reduce_over_shards:
-      lp = distribute_lib.psum(lp, self.experimental_shard_axis_names)
-
-    return lp
+  def lp_fn(self, x):
+    lp = getattr(self.distribution, fn_name)(x)
+    return distribute_lib.psum(lp, self.experimental_shard_axis_names)
 
   lp_fn.__name__ = f'_{fn_name}'
   return lp_fn
@@ -181,7 +174,7 @@ def _default_event_space_bijector(self, *args, **kwargs):
 
 
 @log_prob_ratio.RegisterLogProbRatio(Sharded)
-def _sharded_log_prob_ratio(p, x, q, y, name=None, reduce_over_shards=True):
+def _sharded_log_prob_ratio(p, x, q, y, name=None):
   """Distributed log-prob ratio for Sharded."""
   with tf.name_scope(name or 'sharded_log_prob_ratio'):
     if p.experimental_shard_axis_names != q.experimental_shard_axis_names:
@@ -194,10 +187,7 @@ def log_prob_ratio_fn(x, y):
       return log_prob_ratio.log_prob_ratio(p.distribution, x,
                                            q.distribution, y)
 
-    if reduce_over_shards:
-      axes = p.experimental_shard_axis_names
-
-      return distribute_lib.make_psum_function(
-          log_prob_ratio_fn, in_axes=(axes, axes), out_axes=axes,
-          out_dtype=x)(x, y)
-    return log_prob_ratio_fn(x, y)
+    axes = p.experimental_shard_axis_names
+    return distribute_lib.make_psum_function(
+        log_prob_ratio_fn, in_axes=(axes, axes), out_axes=axes,
+        out_dtype=x)(x, y)