More experiments.

Author: jejjohnson

jupyterbook/content/experimental/cifar/gf_cifar10_naive_tfds_splines.ipynb

@@ -11,8 +11,8 @@
 from distrax._src.distributions.normal import Normal
 
 from rbig_jax.transforms.base import Bijector, BijectorChain
-from rbig_jax.transforms.inversecdf import InitInverseGaussCDF
-from rbig_jax.transforms.logit import InitLogitTransform
+from rbig_jax.transforms.inversecdf import InitInverseGaussCDF, InitGaussCDF
+from rbig_jax.transforms.logit import InitLogitTransform, InitSigmoidTransform
 from rbig_jax.transforms.parametric.householder import InitHouseHolder
 from rbig_jax.transforms.parametric.mixture.gaussian import InitMixtureGaussianCDF
 from rbig_jax.transforms.parametric.mixture.logistic import InitMixtureLogisticCDF
@@ -253,6 +253,110 @@ def init_gf_spline_model(
     return gf_model
 
 
+def init_gf_composite_spline_model(
+    shape: tuple,
+    X: Array = None,
+    n_blocks: int = 4,
+    n_bins: int = 20,
+    range_min: float = 0.0,
+    range_max: float = 1.0,
+    init_rotation: str = "random",
+    n_reflections: int = 10,
+    squeeze: str = "sigmoid",
+    **kwargs,
+):
+
+    n_features = shape[0]
+    rng = jax.random.PRNGKey(42)
+    # rng, _ = jax.random.split(jax.random.PRNGKey(123), 2)
+    # =====================
+    # Composite Transform
+    # ======================
+    init_nl_forward_f = InitGaussCDF()
+    init_nl_inverse_f = InitInverseGaussCDF()
+
+    # =====================
+    # RQ Spline
+    # ======================
+    init_rq_f = InitPiecewiseRationalQuadraticCDF(
+        n_bins=n_bins, range_min=range_min, range_max=range_max, **kwargs
+    )
+    # =====================
+    # HouseHolder Transform
+    # ======================
+    n_reflections = n_reflections
+    # initialize init function
+    init_hh_f = InitHouseHolder(n_reflections=n_reflections, method=init_rotation)
+
+    block_rngs = jax.random.split(rng, num=n_blocks)
+    # rng = jax.random.split(jax.random.PRNGKey(42), n_blocks)
+    # block_rngs = jax.random.split(jax.random.PRNGKey(42), n_blocks)
+
+    itercount = itertools.count()
+    bijectors = []
+
+    X_g = X.copy()
+
+    pbar = tqdm.tqdm(block_rngs)
+    with pbar:
+        for iblock, irng in enumerate(pbar):
+
+            pbar.set_description(
+                f"Initializing - Block: {iblock+1} | Layer {next(itercount)}"
+            )
+            # ======================
+            # Forward Squeezing Transform
+            # ======================
+            # intialize bijector and transformation
+            X_g, layer = init_nl_forward_f.transform_and_bijector(inputs=X_g,)
+            # add bijector to list
+            bijectors.append(layer)
+            # ======================
+            # RQ Spline
+            # ======================
+            # create keys for all inits
+            irng, irq_rng = jax.random.split(irng, 2)
+
+            # intialize bijector and transformation
+            X_g, layer = init_rq_f.transform_and_bijector(
+                inputs=X_g, rng=irq_rng, shape=X.shape[1:]
+            )
+
+            # add bijector to list
+            bijectors.append(layer)
+
+            # ======================
+            # Inverse Squeezing Transform
+            # ======================
+            # intialize bijector and transformation
+            X_g, layer = init_nl_inverse_f.transform_and_bijector(inputs=X_g,)
+            # add bijector to list
+            bijectors.append(layer)
+
+            # ======================
+            # HOUSEHOLDER
+            # ======================
+            pbar.set_description(
+                f"Initializing - Block: {iblock+1} | Layer {next(itercount)}"
+            )
+            # create keys for all inits
+            irng, hh_rng = jax.random.split(irng, 2)
+
+            # intialize bijector and transformation
+            X_g, layer = init_hh_f.transform_and_bijector(
+                inputs=X_g, rng=hh_rng, n_features=n_features
+            )
+
+            bijectors.append(layer)
+
+    # create base dist
+    base_dist = Normal(jnp.zeros((n_features,)), jnp.ones((n_features,)))
+
+    # create flow model
+    gf_model = GaussianizationFlow(base_dist=base_dist, bijectors=bijectors)
+    return gf_model
+
+
 def add_gf_model_args(parser):
     # ====================
     # Model Args

jupyterbook/content/experimental/details/marginal_gauss.ipynb

@@ -84,8 +84,6 @@ def plot_info_loss(
         plt.show()
 
 
-
-
 def plot_image_grid(image, image_shape: Optional = None):
 
     fig = plt.figure(figsize=(10.0, 10.0))

jupyterbook/content/experimental/mnist/gf_demo_mnist_naive_tfds_splines.ipynb

@@ -129,31 +129,32 @@ def InitGaussCDF(eps: float = 1e-5, jitted=False):
     else:
         f = bijector.forward
 
-    def init_bijector(inputs, **kwargs):
+    def transform(inputs, **kwargs):
+
+        outputs = f(inputs)
+
+        return outputs
+
+    def bijector(inputs=None, **kwargs):
 
         return GaussCDF(eps=eps)
 
-    def bijector_and_transform(inputs, **kwargs):
+    def transform_and_bijector(inputs, **kwargs):
         outputs = f(inputs)
         return outputs, GaussCDF(eps=eps)
 
-    def transform(inputs, **kwargs):
-        outputs = f(inputs)
-        return outputs
+    def transform_gradient_bijector(inputs, **kwargs):
+        bijector = GaussCDF(eps=eps)
 
-    def params(inputs, **kwargs):
-        return ()
+        outputs, logabsdet = bijector.forward_and_log_det(inputs)
 
-    def params_and_transform(inputs, **kwargs):
-        outputs = f(inputs)
-        return outputs, ()
+        return outputs, logabsdet, bijector
 
     return InitLayersFunctions(
-        bijector=init_bijector,
-        bijector_and_transform=bijector_and_transform,
         transform=transform,
-        params=params,
-        params_and_transform=params_and_transform,
+        bijector=bijector,
+        transform_and_bijector=transform_and_bijector,
+        transform_gradient_bijector=transform_gradient_bijector,
     )