Add local measure to discrete distributions. This ensures Discrete distributions transform correctly under bijectors.

```python
dist = tfp.distributions.Bernoulli(probs=0.5, dtype=tf.float32)
transformed_dist = tfp.bijectors.Scale(2.)(dist)
transformed_dist.prob(0.)  # Expect this to be 0.5, but if we apply a det jacobian correction of 1 / 2, this would be 0.25
```
PiperOrigin-RevId: 452449290

Author: srvasude

tensorflow_probability/python/distributions/BUILD

@@ -309,6 +309,7 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:sigmoid",
+        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:batched_rejection_sampler",
         "//tensorflow_probability/python/internal:distribution_util",
@@ -1565,6 +1566,7 @@ multi_substrate_py_library(
         ":distribution",
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:softmax_centered",
+        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
@@ -1584,6 +1586,7 @@ multi_substrate_py_library(
         ":gamma",
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:sigmoid",
+        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",

tensorflow_probability/python/distributions/bernoulli.py

@@ -28,7 +28,9 @@
 from tensorflow_probability.python.internal import tensor_util
 
 
-class Bernoulli(distribution.AutoCompositeTensorDistribution):
+class Bernoulli(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Bernoulli distribution.
 
   The Bernoulli distribution with `probs` parameter, i.e., the probability of a

tensorflow_probability/python/distributions/beta_binomial.py

@@ -46,7 +46,9 @@
 """
 
 
-class BetaBinomial(distribution.AutoCompositeTensorDistribution):
+class BetaBinomial(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Beta-Binomial compound distribution.
 
   The Beta-Binomial distribution is parameterized by (a batch of) `total_count`

tensorflow_probability/python/distributions/binomial.py

@@ -262,7 +262,9 @@ def _random_binomial(
     return sampler_impl(**params)
 
 
-class Binomial(distribution.AutoCompositeTensorDistribution):
+class Binomial(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Binomial distribution.
 
   This distribution is parameterized by `probs`, a (batch of) probabilities for

tensorflow_probability/python/distributions/categorical.py

@@ -58,7 +58,9 @@ def _broadcast_cat_event_and_params(event, params, base_dtype):
   return event, params
 
 
-class Categorical(distribution.AutoCompositeTensorDistribution):
+class Categorical(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Categorical distribution over integers.
 
   The Categorical distribution is parameterized by either probabilities or

tensorflow_probability/python/distributions/dirichlet_multinomial.py

@@ -51,7 +51,9 @@
 with `self.concentration` and `self.total_count`."""
 
 
-class DirichletMultinomial(distribution.AutoCompositeTensorDistribution):
+class DirichletMultinomial(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Dirichlet-Multinomial compound distribution.
 
   The Dirichlet-Multinomial distribution is parameterized by a (batch of)

tensorflow_probability/python/distributions/distribution.py

@@ -2119,6 +2119,48 @@ class MyDistribution(tfb.AutoCompositeTensorDistribution):
   pass
 
 
+class DiscreteDistributionMixin(object):
+  """Mixin for Distributions over discrete spaces.
+
+  This mixin identifies a `Distribution` as a discrete distribution, which in
+  turn ensures that it is transformed properly under `TransformedDistribution`.
+
+  Normally, for a continuous distribution `dist` by a bijector `bij`, we have
+  the following formula for the `log_prob`:
+      `dist.log_prob(bij.inverse(y)) + bij.inverse_log_det_jacobian(y)`.
+  For a discrete distribution, we don't apply the `inverse_log_det_jacobian`
+  correction (hence just `dist.log_prob(bij.inverse(y))`). This difference
+  comes from transforming a probability density vs. probabilities.
+
+  As an example, we could take a Bernoulli distribution (
+  whose samples are `0` or `1`) and square it via `tfb.Square`. Samples from
+  this new distribution are still `0` or `1` and one would expect that the
+  probabilities for `0` and `1` are unchanged after this transformation.
+
+  ```python
+  dist = tfp.distributions.Bernoulli(probs=0.5)
+  dist.prob(1.)  # expect 0.5
+  transformed_dist = tfp.bijectors.Square()(dist)
+  transformed_dist.prob(1.) # expect 0.5
+  ```
+
+  If we apply the jacobian correction, we would instead get the wrong answer
+
+  ```python
+  # If we compute with the jacobian correction explicitly, we get the wrong
+  # answer.
+  bij = tfp.bijectors.Square()
+  prob_at_1 = dist.log_prob(bij.inverse(1.)) + bij.inverse_log_det_jacobian(1.)
+  prob_at_1 = tf.math.exp(prob_at_1) # This is 0.25
+  ```
+  """
+
+  @property
+  def _experimental_tangent_space(self):
+    from tensorflow_probability.python.experimental import tangent_spaces  # pylint: disable=g-import-not-at-top
+    return tangent_spaces.ZeroSpace()
+
+
 class _PrettyDict(dict):
   """`dict` with stable `repr`, `str`."""
 

tensorflow_probability/python/distributions/distribution_properties_test.py

@@ -91,6 +91,12 @@
 })
 
 
+DISCRETE_BUT_NOT_TRANSFORMABLE = [
+    # Samples are integers so can't be transformed by a float bijector.
+    'DeterminantalPointProcess',
+]
+
+
 @test_util.test_all_tf_execution_regimes
 class StatisticConsistentShapesTest(test_util.TestCase):
 
@@ -230,6 +236,43 @@ def testDistribution(self, dist_name, data):
     self.assertAllEqual(s1, s2)
 
 
+@test_util.test_all_tf_execution_regimes
+class TestDiscreteDistributions(test_util.TestCase):
+
+  @parameterized.named_parameters(
+      {'testcase_name': dname, 'dist_name': dname}
+      for dname in sorted(list(set(dhps.DISCRETE_DISTS) -
+                               set(DISCRETE_BUT_NOT_TRANSFORMABLE))))
+  @hp.given(hps.data())
+  @tfp_hps.tfp_hp_settings()
+  def testNoJacobianCorrection(self, dist_name, data):
+
+    # Disable validate args since transforming with Softplus and inverting
+    # might make arguments not as close to integers.
+    dist = data.draw(dhps.distributions(
+        dist_name=dist_name,
+        enable_vars=False,
+        validate_args=False))
+
+    # Ensure that these are distributions over floats so we can apply the
+    # Softplus bijector.
+    if 'dtype' in dist.parameters:
+      dist = dist.copy(dtype=tf.float32)
+    bij = tfb.Softplus()
+    transformed_dist = tfd.TransformedDistribution(dist, bijector=bij)
+
+    seed = test_util.test_seed()
+    samples = transformed_dist.sample(7, seed=seed)
+    # Break bijector caching.
+    samples = self.evaluate(
+        samples + tf.constant(0., dtype=samples.dtype))
+
+    # Check that no jacobian correction is added for a discrete distribution.
+    self.assertAllClose(
+        self.evaluate(dist.log_prob(bij.inverse(samples))),
+        self.evaluate(transformed_dist.log_prob(samples)))
+
+
 @test_util.test_all_tf_execution_regimes
 class SampleAndLogProbTest(test_util.TestCase):
 
@@ -630,8 +673,8 @@ class TestMixingGraphAndEagerModes(test_util.TestCase):
 
   @parameterized.named_parameters(
       {'testcase_name': dname, 'dist_name': dname}
-      for dname in  sorted(list(dhps.INSTANTIABLE_BASE_DISTS.keys()) +
-                           list(dhps.INSTANTIABLE_META_DISTS))
+      for dname in sorted(list(dhps.INSTANTIABLE_BASE_DISTS.keys()) +
+                          list(dhps.INSTANTIABLE_META_DISTS))
   )
   @hp.given(hps.data())
   @tfp_hps.tfp_hp_settings()

tensorflow_probability/python/distributions/dpp.py

@@ -240,7 +240,9 @@ def body(i, vecs, cur_sample, seed):
     return tf.cast(sample, tf.int32)
 
 
-class DeterminantalPointProcess(distribution.AutoCompositeTensorDistribution):
+class DeterminantalPointProcess(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Determinantal point process (DPP) distribution.
 
   The DPP disribution parameterized by the eigenvalues and eigenvectors of the

tensorflow_probability/python/distributions/empirical.py

@@ -52,7 +52,9 @@ def _broadcast_event_and_samples(event, samples, event_ndims):
   return event, samples
 
 
-class Empirical(distribution.AutoCompositeTensorDistribution):
+class Empirical(
+    distribution.DiscreteDistributionMixin,
+    distribution.AutoCompositeTensorDistribution):
   """Empirical distribution.
 
   The Empirical distribution is parameterized by a [batch] multiset of samples.