Cascading flow surrogate posterior

tensorflow_probability/python/experimental/vi/BUILD

@@ -31,6 +31,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":automatic_structured_vi",
+        ":cascading_flows",
         ":surrogate_posteriors",
         "//tensorflow_probability/python/experimental/vi/util",
         "//tensorflow_probability/python/internal:all_util",
@@ -67,6 +68,37 @@ py_library(
     ],
 )
 
+py_library(
+    name = "cascading_flows",
+    srcs = ["cascading_flows.py.py"],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:build_highway_flow_layer",
+        "//tensorflow_probability/python/bijectors:chain",
+        "//tensorflow_probability/python/bijectors:reshape",
+        "//tensorflow_probability/python/bijectors:scale",
+        "//tensorflow_probability/python/bijectors:shift",
+        "//tensorflow_probability/python/bijectors:split",
+        "//tensorflow_probability/python/distributions:batch_broadcast",
+        "//tensorflow_probability/python/distributions:beta",
+        "//tensorflow_probability/python/distributions:blockwise",
+        "//tensorflow_probability/python/distributions:chi2",
+        "//tensorflow_probability/python/distributions:exponential",
+        "//tensorflow_probability/python/distributions:gamma",
+        "//tensorflow_probability/python/distributions:half_normal",
+        "//tensorflow_probability/python/distributions:joint_distribution_auto_batched",
+        "//tensorflow_probability/python/distributions:joint_distribution_coroutine",
+        "//tensorflow_probability/python/distributions:normal",
+        "//tensorflow_probability/python/distributions:sample",
+        "//tensorflow_probability/python/distributions:transformed_distribution",
+        "//tensorflow_probability/python/distributions:truncated_normal",
+        "//tensorflow_probability/python/distributions:uniform",
+        "//tensorflow_probability/python/experimental/bijectors:build_trainable_highway_flow",
+        "//tensorflow_probability/python/internal:samplers",
+    ],
+)
+
 py_library(
     name = "surrogate_posteriors",
     srcs = ["surrogate_posteriors.py"],
@@ -111,6 +143,22 @@ py_test(
     ],
 )
 
+py_test(
+    name = "cascading_flows_test",
+    size = "large",
+    srcs = ["cascading_flows_test.py"],
+    python_version = "PY3",
+    shard_count = 4,
+    srcs_version = "PY3",
+    deps = [
+        # absl/testing:parameterized dep,
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
 py_test(
     name = "surrogate_posteriors_test",
     size = "large",

tensorflow_probability/python/experimental/vi/__init__.py

@@ -17,6 +17,7 @@
 from tensorflow_probability.python.experimental.vi import util
 from tensorflow_probability.python.experimental.vi.automatic_structured_vi import build_asvi_surrogate_posterior
 from tensorflow_probability.python.experimental.vi.automatic_structured_vi import register_asvi_substitution_rule
+from tensorflow_probability.python.experimental.vi.cascading_flows import build_cf_surrogate_posterior
 from tensorflow_probability.python.experimental.vi.surrogate_posteriors import build_affine_surrogate_posterior
 from tensorflow_probability.python.experimental.vi.surrogate_posteriors import build_affine_surrogate_posterior_from_base_distribution
 from tensorflow_probability.python.experimental.vi.surrogate_posteriors import build_factored_surrogate_posterior
@@ -29,6 +30,7 @@
     'build_affine_surrogate_posterior',
     'build_affine_surrogate_posterior_from_base_distribution',
     'build_asvi_surrogate_posterior',
+    'build_cf_surrogate_posterior',
     'build_factored_surrogate_posterior',
     'build_split_flow_surrogate_posterior',
     'build_trainable_location_scale_distribution',