Enable AutoCompositeTensor for simple distributions.

PiperOrigin-RevId: 378226943

Author: emilyfertig

tensorflow_probability/python/distributions/BUILD

@@ -384,6 +384,7 @@ multi_substrate_py_library(
     srcs = ["chi.py"],
     deps = [
         ":chi2",
+        ":distribution",
         ":kullback_leibler",
         ":transformed_distribution",
         # numpy dep,
@@ -610,8 +611,11 @@ multi_substrate_py_library(
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:identity",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
+        "//tensorflow_probability/python/internal:prefer_static",
+        "//tensorflow_probability/python/internal:reparameterization",
         "//tensorflow_probability/python/internal:samplers",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
@@ -806,6 +810,7 @@ multi_substrate_py_library(
     name = "gumbel",
     srcs = ["gumbel.py"],
     deps = [
+        ":distribution",
         ":kullback_leibler",
         ":transformed_distribution",
         ":uniform",
@@ -826,6 +831,7 @@ multi_substrate_py_library(
     name = "gev",
     srcs = ["gev.py"],
     deps = [
+        ":distribution",
         ":kullback_leibler",
         ":transformed_distribution",
         ":uniform",
@@ -995,6 +1001,7 @@ multi_substrate_py_library(
     name = "johnson_su",
     srcs = ["johnson_su.py"],
     deps = [
+        ":distribution",
         ":normal",
         ":transformed_distribution",
         # tensorflow dep,
@@ -1132,6 +1139,7 @@ multi_substrate_py_library(
     name = "kumaraswamy",
     srcs = ["kumaraswamy.py"],
     deps = [
+        ":distribution",
         ":transformed_distribution",
         ":uniform",
         # numpy dep,
@@ -1278,6 +1286,7 @@ multi_substrate_py_library(
     name = "lognormal",
     srcs = ["lognormal.py"],
     deps = [
+        ":distribution",
         ":normal",
         ":transformed_distribution",
         # numpy dep,
@@ -1294,6 +1303,7 @@ multi_substrate_py_library(
     name = "logitnormal",
     srcs = ["logitnormal.py"],
     deps = [
+        ":distribution",
         ":normal",
         ":transformed_distribution",
         # numpy dep,
@@ -1426,6 +1436,7 @@ multi_substrate_py_library(
     name = "moyal",
     srcs = ["moyal.py"],
     deps = [
+        ":distribution",
         ":kullback_leibler",
         ":transformed_distribution",
         ":uniform",
@@ -1500,6 +1511,7 @@ multi_substrate_py_library(
     name = "mvn_linear_operator",
     srcs = ["mvn_linear_operator.py"],
     deps = [
+        ":distribution",
         ":kullback_leibler",
         ":normal",
         ":sample",
@@ -2219,6 +2231,7 @@ multi_substrate_py_library(
     name = "vector_exponential_linear_operator",
     srcs = ["vector_exponential_linear_operator.py"],
     deps = [
+        ":distribution",
         ":exponential",
         ":sample",
         ":transformed_distribution",
@@ -2227,6 +2240,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:scale_matvec_linear_operator",
         "//tensorflow_probability/python/bijectors:shift",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:tensorshape_util",
@@ -2287,6 +2301,7 @@ multi_substrate_py_library(
     srcs = ["weibull.py"],
     deps = [
         ":distribution",
+        ":transformed_distribution",
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:chain",
@@ -3382,6 +3397,7 @@ multi_substrate_py_test(
     name = "mvn_diag_plus_low_rank_covariance_test",
     srcs = ["mvn_diag_plus_low_rank_covariance_test.py"],
     deps = [
+        ":distribution",
         ":mvn_diag_plus_low_rank_covariance",
         # numpy dep,
         # tensorflow dep,

tensorflow_probability/python/distributions/bates.py

@@ -58,7 +58,7 @@
 }
 
 
-class Bates(distribution.Distribution):
+class Bates(distribution.AutoCompositeTensorDistribution):
   """Bates distribution.
 
   The Bates distribution is the distribution of the average of `total_count`

tensorflow_probability/python/distributions/bernoulli.py

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

tensorflow_probability/python/distributions/beta.py

@@ -48,7 +48,7 @@
 `[0, 1].` It must have a shape compatible with `self.batch_shape()`."""
 
 
-class Beta(distribution.Distribution):
+class Beta(distribution.AutoCompositeTensorDistribution):
   """Beta distribution.
 
   The Beta distribution is defined over the `(0, 1)` interval using parameters

tensorflow_probability/python/distributions/beta_binomial.py

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

tensorflow_probability/python/distributions/beta_quotient.py

@@ -41,7 +41,7 @@
 ]
 
 
-class BetaQuotient(distribution.Distribution):
+class BetaQuotient(distribution.AutoCompositeTensorDistribution):
   """BetaQuotient distribution.
 
   The Beta Quotient distribution is defined over the positive reals, as

tensorflow_probability/python/distributions/binomial.py

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

tensorflow_probability/python/distributions/categorical.py

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

tensorflow_probability/python/distributions/cauchy.py

@@ -40,7 +40,7 @@
 ]
 
 
-class Cauchy(distribution.Distribution):
+class Cauchy(distribution.AutoCompositeTensorDistribution):
   """The Cauchy distribution with location `loc` and scale `scale`.
 
   #### Mathematical details

tensorflow_probability/python/distributions/chi.py

@@ -26,6 +26,7 @@
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.bijectors import square as square_bijector
 from tensorflow_probability.python.distributions import chi2
+from tensorflow_probability.python.distributions import distribution
 from tensorflow_probability.python.distributions import kullback_leibler
 from tensorflow_probability.python.distributions import transformed_distribution
 from tensorflow_probability.python.internal import assert_util
@@ -34,7 +35,10 @@
 from tensorflow_probability.python.internal import tensor_util
 
 
-class Chi(transformed_distribution.TransformedDistribution):
+# TODO(b/182603117): Remove `AutoCompositeTensor` subclass when
+# `TransformedDistribution` is converted to `CompositeTensor`.
+class Chi(transformed_distribution.TransformedDistribution,
+          distribution.AutoCompositeTensorDistribution):
   """Chi distribution.
 
   The Chi distribution is defined over nonnegative real numbers and uses a