makes it easier to change base distribution

src/bernstein_flow/distributions/bernstein_flow.py

@@ -161,6 +161,39 @@ def init_bijectors(
         return tfb.Invert(tfb.Chain(bijectors))
 
 
+def get_base_distribution(base_distribution, dtype, **kwds):
+    if isinstance(base_distribution, tfd.Distribution):
+        return base_distribution
+    else:
+        if base_distribution == "normal":
+            default_kwds = dict(loc=tf.convert_to_tensor(0, dtype=dtype), scale=1.0)
+            default_kwds.update(**kwds)
+            dist = tfd.Normal(**default_kwds)
+        elif base_distribution == "truncated_normal":
+            default_kwds = dict(
+                loc=tf.convert_to_tensor(0, dtype=dtype), scale=1.0, low=-4, high=4
+            )
+            default_kwds.update(**kwds)
+            dist = tfd.TruncatedNormal(**default_kwds)
+        elif base_distribution == "log_normal":
+            default_kwds = dict(loc=tf.convert_to_tensor(0, dtype=dtype), scale=1.0)
+            default_kwds.update(**kwds)
+            dist = tfd.LogNormal(**default_kwds)
+        elif base_distribution == "logistic":
+            default_kwds = dict(loc=tf.convert_to_tensor(0, dtype=dtype), scale=1.0)
+            default_kwds.update(**kwds)
+            dist = tfd.Logistic(**default_kwds)
+        elif base_distribution == "uniform":
+            default_kwds = dict(low=tf.convert_to_tensor(0, dtype=dtype), high=1.0)
+            default_kwds.update(**kwds)
+            dist = tfd.Uniform(**default_kwds)
+        elif base_distribution == "kumaraswamy":
+            dist = tfd.Kumaraswamy(**kwds)
+        else:
+            raise ValueError(f"Unsupported distribution type {base_distribution}.")
+    return dist
+
+
 class BernsteinFlow(tfd.TransformedDistribution):
     """
     This class implements a `tfd.TransformedDistribution` using Bernstein
@@ -174,6 +207,7 @@ def __init__(
         b1=None,
         a2=None,
         base_distribution=None,
+        base_distribution_kwds={},
         clip_to_bernstein_domain=True,
         clip_base_distribution=False,
         bb_class=BernsteinBijector,
@@ -199,9 +233,16 @@ def __init__(
             shape = prefer_static.shape(thetas)
 
             if base_distribution is None:
-                base_distribution = tfd.Normal(
-                    loc=tf.zeros(shape[:-1], dtype=dtype), scale=1.0
-                )
+                if bb_class == BernsteinBijector:
+                    base_distribution = get_base_distribution(
+                        "truncated_normal", dtype, **base_distribution_kwds
+                    )
+                else:
+                    base_distribution = get_base_distribution(
+                        "normal", dtype, **base_distribution_kwds
+                    )
+            else:
+                base_distribution = get_base_distribution(base_distribution, dtype)
 
             bijector = init_bijectors(
                 thetas,

test/distributions/test_bernstein_flow.py

@@ -56,15 +56,15 @@ def gen_dist(batch_shape, order=5, dtype=tf.float32, seed=1, **kwds):
 
 
 class BernsteinFlowTest(tf.test.TestCase):
-    def f(self, normal_dist, trans_dist):
+    def f(self, normal_dist, trans_dist, stay_in_domain=False):
         tf.random.set_seed(42)
 
         dtype = normal_dist.dtype
         for input_shape in [[1], [1, 1], [1] + normal_dist.batch_shape]:
             x = tf.random.uniform(
                 shape=[10] + normal_dist.batch_shape,
-                minval=-100,
-                maxval=100,
+                minval=0 if stay_in_domain else -100,
+                maxval=1 if stay_in_domain else 100,
                 dtype=dtype,
             )
 
@@ -86,12 +86,12 @@ def f(self, normal_dist, trans_dist):
             # check the Normalization
             self.assertAllClose(
                 normal_dist.cdf(normal_dist.dtype.min),
-                trans_dist.cdf(trans_dist.dtype.min),
+                trans_dist.cdf(0.0 if stay_in_domain else trans_dist.dtype.min),
                 atol=1e-3,
             )
             self.assertAllClose(
                 normal_dist.cdf(normal_dist.dtype.max),
-                trans_dist.cdf(trans_dist.dtype.max),
+                trans_dist.cdf(1.0 if stay_in_domain else trans_dist.dtype.max),
                 atol=1e-3,
             )
 
@@ -153,14 +153,13 @@ def test_dist_multi_extra(self):
     @pytest.mark.skip
     def test_log_normal(self):
         batch_shape = [16, 10]
-        log_normal = tfd.LogNormal(loc=tf.zeros(batch_shape, dtype=dtype), scale=1.0)
         normal_dist, trans_dist = gen_dist(
             batch_shape=batch_shape,
             order=10,
             dtype=dtype,
-            base_distribution=log_normal,
+            base_distribution="log_normal",
             thetas_constrain_fn=get_thetas_constrain_fn(
-                low=1e-10, high=tf.math.exp(tf.constant(4.0, dtype=dtype))
+                low=1e-12, high=tf.math.exp(tf.constant(6.0, dtype=dtype))
             ),
             scale_base_distribution=True,
         )
@@ -169,34 +168,55 @@ def test_log_normal(self):
     @pytest.mark.skip
     def test_logistic(self):
         batch_shape = [16, 10]
-        logistic = tfd.Logistic(loc=tf.zeros(batch_shape, dtype=dtype), scale=1)
         normal_dist, trans_dist = gen_dist(
             batch_shape=batch_shape,
             order=10,
             dtype=dtype,
-            base_distribution=logistic,
+            base_distribution="logistic",
+            bb_class=BernsteinBijectorLinearExtrapolate,
             thetas_constrain_fn=get_thetas_constrain_fn(
-                low=-8, high=8, allow_flexible_bounds=True
+                low=-20, high=20, allow_flexible_bounds=True
             ),
-            scale_base_distribution=True,
+            scale_base_distribution=False,
+            clip_to_bernstein_domain=False,
         )
         self.f(normal_dist, trans_dist)
 
     @pytest.mark.skip
     def test_uniform(self):
         batch_shape = [16, 10]
-        uniform = tfd.Uniform(
-            -tf.ones(batch_shape, dtype=dtype), tf.ones(batch_shape, dtype=dtype)
+        normal_dist, trans_dist = gen_dist(
+            batch_shape=batch_shape,
+            order=10,
+            dtype=dtype,
+            base_distribution="uniform",
+            thetas_constrain_fn=get_thetas_constrain_fn(low=0.0, high=1.0),
+            clip_to_bernstein_domain=False,
+            scale_base_distribution=False,
+            shift_data=False,
+            scale_data=False,
         )
+        self.f(normal_dist, trans_dist, stay_in_domain=True)
+
+    @pytest.mark.skip
+    def test_kumaraswamy(self):
+        batch_shape = [16, 10]
         normal_dist, trans_dist = gen_dist(
             batch_shape=batch_shape,
             order=10,
             dtype=dtype,
-            base_distribution=uniform,
-            thetas_constrain_fn=get_thetas_constrain_fn(low=-1.0, high=1.0),
+            base_distribution="kumaraswamy",
+            base_distribution_kwds={
+                "concentration1": tf.convert_to_tensor(5.0, dtype),
+                "concentration0": 2.0,
+            },
+            thetas_constrain_fn=get_thetas_constrain_fn(low=0.0, high=1.0),
+            clip_to_bernstein_domain=False,
             scale_base_distribution=False,
+            shift_data=False,
+            scale_data=False,
         )
-        self.f(normal_dist, trans_dist)
+        self.f(normal_dist, trans_dist, stay_in_domain=True)
 
     @pytest.mark.skip
     def test_student_t(self):