Add MarkovChain meta distribution for time series modeling.

My original motivation for this was to unify the two stochastic volatility models in the inference gym (https://github.com/tensorflow/probability/blob/master/spinoffs/inference_gym/inference_gym/targets/stochastic_volatility.py and
https://github.com/tensorflow/probability/blob/master/spinoffs/inference_gym/inference_gym/targets/vectorized_stochastic_volatility.py). Currently, we have a non-vectorized model that exposes the graphical model structure needed for ASVI, but is very slow, and a vectorized version that is fast but does not expose any structure. Representing the volatility series as a MarkovChain would expose the graphical structure *and* support vectorized log-prob evaluation.

More generally we could think of MarkovChain as a new flavor of joint distribution, since (like TransformedDistribution) it can represent multipart events. It proposes an answer to the question, "How can a JD efficiently define random variables in a loop?" Currently the only way to do this is to write an explicit Python loop in a JDCoroutine model, as the naive stochastic volatility model does, but this is obviously suboptimal. MarkovChains might be useful both as joint distributions over time series and to represent time-series components in larger joint distribution models.

Another potential use is as a bridge between the particle filter API and the Distributions API. Our particle filtering code expects a time series model to be specified in terms of an initial prior, a transition function, and an observation function. It's natural to want to represent the distribution on the latent sequences implied by the prior and transition model, e.g., for prior predictive checks; MarkovChain allows us to do this.

I'm not totally sure if MarkovChain is the right name. Throwing out some other possibilities: Markov, AutoregressiveMarkov, TimeSeries, Loop, Scan, <your suggestion here>?

PiperOrigin-RevId: 377150791

Author: davmre

tensorflow_probability/python/distributions/BUILD

@@ -105,6 +105,7 @@ multi_substrate_py_library(
         ":logitnormal",
         ":loglogistic",
         ":lognormal",
+        ":markov_chain",
         ":masked",
         ":matrix_normal_linear_operator",
         ":matrix_t_linear_operator",
@@ -1305,6 +1306,27 @@ multi_substrate_py_library(
     ],
 )
 
+multi_substrate_py_library(
+    name = "markov_chain",
+    srcs = ["markov_chain.py"],
+    deps = [
+        ":distribution",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:chain",
+        "//tensorflow_probability/python/bijectors:invert",
+        "//tensorflow_probability/python/bijectors:joint_map",
+        "//tensorflow_probability/python/bijectors:split",
+        "//tensorflow_probability/python/internal:assert_util",
+        "//tensorflow_probability/python/internal:distribution_util",
+        "//tensorflow_probability/python/internal:parameter_properties",
+        "//tensorflow_probability/python/internal:prefer_static",
+        "//tensorflow_probability/python/internal:samplers",
+        "//tensorflow_probability/python/internal:tensor_util",
+        "//tensorflow_probability/python/internal:tensorshape_util",
+    ],
+)
+
 multi_substrate_py_library(
     name = "matrix_normal_linear_operator",
     srcs = ["matrix_normal_linear_operator.py"],
@@ -3218,6 +3240,20 @@ multi_substrate_py_test(
     ],
 )
 
+multi_substrate_py_test(
+    name = "markov_chain_test",
+    size = "medium",
+    srcs = ["markov_chain_test.py"],
+    deps = [
+        # numpy dep,
+        # scipy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/distributions",
+        "//tensorflow_probability/python/internal:samplers",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
 multi_substrate_py_test(
     name = "matrix_normal_linear_operator_test",
     size = "medium",

tensorflow_probability/python/distributions/__init__.py

@@ -85,6 +85,7 @@
 from tensorflow_probability.python.distributions.loglogistic import LogLogistic
 from tensorflow_probability.python.distributions.lognormal import LogNormal
 from tensorflow_probability.python.distributions.logitnormal import LogitNormal
+from tensorflow_probability.python.distributions.markov_chain import MarkovChain
 from tensorflow_probability.python.distributions.masked import Masked
 from tensorflow_probability.python.distributions.matrix_normal_linear_operator import MatrixNormalLinearOperator
 from tensorflow_probability.python.distributions.matrix_t_linear_operator import MatrixTLinearOperator
@@ -230,6 +231,7 @@
     'LogLogistic',
     'LogNormal',
     'LogitNormal',
+    'MarkovChain',
     'Moyal',
     'NegativeBinomial',
     'Normal',

tensorflow_probability/python/distributions/hypothesis_testlib.py

@@ -80,6 +80,7 @@
     'LambertWDistribution',
     'MatrixNormalLinearOperator',
     'MatrixTLinearOperator',
+    'MarkovChain',
     'Masked',  # (has strategy)
     'Mixture',  # (has strategy)
     'MixtureSameFamily',  # (has strategy)