Add more surjectors

dirmeier · dirmeier · Nov 28, 2022 · Nov 4, 2022 · Nov 5, 2022 · Nov 11, 2022
commit ae0bb62c04dde0fe8ae25e64364162b2e4223e0c
diff --git a/README.md b/README.md
@@ -1,39 +1,37 @@
 # surjectors
 
 [![status](http://www.repostatus.org/badges/latest/concept.svg)](http://www.repostatus.org/#concept)
-[![ci](https://github.com/dirmeier/surjectors/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/Inference surjection layers/actions/workflows/ci.yaml)
+[![ci](https://github.com/dirmeier/surjectors/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/surjectors/actions/workflows/ci.yaml)
 [![codecov](https://codecov.io/gh/dirmeier/surjectors/branch/main/graph/badge.svg)](https://codecov.io/gh/dirmeier/surjectors)
 [![codacy]()]()
 [![documentation](https://readthedocs.org/projects/surjectors/badge/?version=latest)](https://surjectors.readthedocs.io/en/latest/?badge=latest)
 [![version](https://img.shields.io/pypi/v/surjectors.svg?colorB=black&style=flat)](https://pypi.org/project/surjectors/)
 
-> Inference surjection layers
+> Surjection layers for density estimation with normalizing flows
 
 ## About
 
-TODO
+Surjectors is a light-weight library of inference and generative surjection layers, i.e., layers that reduce dimensionality, for density estimation using normalizing flows.
+Surjectors builds on Distrax and Haiku.
 
 ## Example usage
 
 TODO
 
 ## Installation
 
-
-To install the latest GitHub <RELEASE>, just call the following on the
-command line:
+To install the latest GitHub <RELEASE>, just call the following on the command line:
 
 ```bash
-pip install git+https://github.com/dirmeier/Inference surjection layers@<RELEASE>
+pip install git+https://github.com/dirmeier/surjectors@<RELEASE>
 ```
 
 ## Contributing
 
 In order to contribute:
 
-1) Fork and download the repository,
-2) create a branch with the name of your new feature (something like `issue/fix-bug-related-to-something` or `feature/implement-new-bound`),
-3) install `surjectors` and dev dependencies via `poetry install` (you might need to create a new `conda` or `venv` environment, to not break other dependencies),
+1) Fork and download the forked repository,
+2) create a branch with the name of your new feature (something like `issue/fix-bug-related-to-something` or `feature/implement-new-surjector`),
+3) install `surjectors` and dev dependencies via `poetry install` (you might want to create a new `conda` or `venv` environment, to not break other dependencies),
 4) develop code, commit changes and push it to your branch,
 5) create a PR
-
diff --git a/surjectors/data.py → examples/data.py b/surjectors/data.py → examples/data.py
diff --git a/examples/multivariate_gaussian.py b/examples/multivariate_gaussian.py
@@ -0,0 +1,32 @@
+prng = hk.PRNGSequence(jax.random.PRNGKey(42))
+matrix = jax.random.uniform(next(prng), (4, 4))
+bias = jax.random.normal(next(prng),  (4,))
+bijector = LowerUpperTriangularAffine(matrix, bias)
+
+#
+# def loss():
+#     x, lc = bijector.inverse_and_log_det(jnp.zeros(4) * 2.1)
+#     lp = distrax.Normal(jnp.zeros(4)).log_prob(x)
+#     return -jnp.sum(lp - lc)
+#
+# print(bijector.matrix)
+#
+# adam = optax.adam(0.003)
+# g = jax.grad(loss)()
+#
+# print(g)
+#
+
+matrix = jax.random.uniform(next(prng), (4, 4))
+bias = jax.random.normal(next(prng),  (4,))
+bijector = LowerUpperTriangularAffine(matrix, bias)
+
+n = jnp.ones((4, 4)) * 3.1
+n += jnp.triu(n) * 2
+
+bijector = LowerUpperTriangularAffine(n, jnp.zeros(4))
+
+
+print(bijector.forward(jnp.ones(4)))
+
+print(n @jnp.ones(4) )
diff --git a/examples/solar_dynamo.py b/examples/solar_dynamo.py
@@ -15,6 +15,3 @@
         jnp.array([549229066, 500358972], dtype=jnp.uint32), 100
     )
     pns[i] = pn
-
-
-Distribution
diff --git a/surjectors/bijectors/__init__.py b/surjectors/bijectors/__init__.py
diff --git a/surjectors/bijectors/coupling.py b/surjectors/bijectors/coupling.py
diff --git a/surjectors/conditioners/__init__.py b/surjectors/conditioners/__init__.py
diff --git a/surjectors/conditioners/transformer.py b/surjectors/conditioners/transformer.py
diff --git a/surjectors/distributions/conditional_distribution.py b/surjectors/distributions/conditional_distribution.py
@@ -0,0 +1,2 @@
+class ConditionalDistribution:
+    pass
diff --git a/surjectors/distribution.py → surjectors/distributions/distribution.py b/surjectors/distribution.py → surjectors/distributions/distribution.py
diff --git a/surjectors/distributions/transformed_distribution.py b/surjectors/distributions/transformed_distribution.py
@@ -1,30 +1,22 @@
-from typing import Tuple
-
+import chex
 import jax
 import jax.numpy as jnp
-from distrax._src.bijectors import bijector as bjct_base
-from distrax._src.distributions import distribution as dist_base
-from distrax._src.utils import conversion
-from tensorflow_probability.substrates import jax as tfp
-
-tfd = tfp.distributions
-
-PRNGKey = dist_base.PRNGKey
-Array = dist_base.Array
-DistributionLike = dist_base.DistributionLike
-BijectorLike = bjct_base.BijectorLike
+from chex import PRNGKey
+from distrax import Distribution
+Array = chex.Array
+from surjectors.surjectors.surjector import Surjector
 
 
 class TransformedDistribution:
-    def __init__(self, base_distribution, surjector):
+    def __init__(self, base_distribution: Distribution, surjector: Surjector):
         self.base_distribution = base_distribution
         self.surjector = surjector
 
-    def log_prob(self, y: Array) -> Array:
-        x, ildj_y = self.surjector.inverse_and_log_det(y)
+    def log_prob(self, y: Array) -> jnp.ndarray:
+        x, lc = self.surjector.inverse_and_likelihood_contribution(y)
         lp_x = self.base_distribution.log_prob(x)
-        lp_y = lp_x + ildj_y
-        return lp_y
+        lp = lp_x - lc
+        return lp
 
     def sample(self, key: PRNGKey, sample_shape=(1,)):
         z = self.base_distribution.sample(seed=key, sample_shape=sample_shape)
@@ -35,6 +27,6 @@ def sample_and_log_prob(self, key: PRNGKey, sample_shape=(1,)):
         z, lp_z = self.base_distribution.sample_and_log_prob(
             seed=key, sample_shape=sample_shape
         )
-        y, fldj = jax.vmap(self.surjector.forward_and_log_det)(z)
-        lp_y = jax.vmap(jnp.subtract)(lp_z, fldj)
-        return y, lp_y
+        y, fldj = jax.vmap(self.surjector.forward_and_likelihood_contribution)(z)
+        lp = jax.vmap(jnp.subtract)(lp_z, fldj)
+        return y, lp
diff --git a/surjectors/surjectors/transform.py → surjectors/surjectors/_transform.py b/surjectors/surjectors/transform.py → surjectors/surjectors/_transform.py
diff --git a/surjectors/surjectors/affine_coupling_funnel.py b/surjectors/surjectors/affine_coupling_funnel.py
@@ -0,0 +1,24 @@
+import chex
+from jax import numpy as jnp
+
+from surjectors.surjectors.surjector import Surjector
+
+
+class AffineCouplingFunnel(Surjector):
+    def __init__(self, n_keep, decoder, transform, encoder, kind="inference_surjection"):
+        super().__init__(n_keep, decoder, encoder, kind)
+        self._transform = transform
+
+    def split_input(self, input):
+        split_proportions = (self.n_keep, input.shape[-1] - self.n_keep)
+        return jnp.split(input, split_proportions, axis=-1)
+
+    def inverse_and_likelihood_contribution(self, y):
+        y_plus, y_minus = self.split_input(y)
+        chex.assert_equal_shape([y_plus, y_minus])
+        z, jac_det = self._transform(y_plus, context=y_minus)
+        lp = self.decoder.log_prob(y_minus, context=z)
+        return z, lp + jac_det
+
+    def forward_and_likelihood_contribution(self, z):
+        raise NotImplementedError()
diff --git a/surjectors/surjectors/chain.py b/surjectors/surjectors/chain.py
@@ -3,6 +3,7 @@
 
 class Chain(Surjector):
     def __init__(self, surjectors):
+        super().__init__(None, None, None, "surjector")
         self._surjectors = surjectors
 
     def inverse_and_likelihood_contribution(self, y):
@@ -18,4 +19,3 @@ def forward_and_likelihood_contribution(self, z):
             x, lc = _surjectors.forward_and_log_det(x)
             log_det += lc
         return y, log_det
-
diff --git a/surjectors/surjectors/funnel.py b/surjectors/surjectors/funnel.py
diff --git a/surjectors/surjectors/linear.py b/surjectors/surjectors/linear.py
diff --git a/surjectors/surjectors/lu_linear.py b/surjectors/surjectors/lu_linear.py
@@ -7,8 +7,11 @@
 
 
 class LULinear(Surjector):
-    def __init__(self, n_keep, dtype=jnp.float32):
+    def __init__(self, n_keep, with_bias=False, dtype=jnp.float32):
         super().__init__(n_keep, None, None, "bijection", dtype)
+        if with_bias:
+            raise NotImplementedError()
+
         n_triangular_entries = ((n_keep - 1) * n_keep) // 2
 
         self._lower_indices = np.tril_indices(n_keep, k=-1)
@@ -51,4 +54,3 @@ def inverse_and_likelihood_contribution(self, y):
 
     def forward_and_likelihood_contribution(self, z):
         pass
-
diff --git a/surjectors/surjectors/mlp.py b/surjectors/surjectors/mlp.py
@@ -0,0 +1,30 @@
+import distrax
+import haiku as hk
+from jax import numpy as jnp
+
+from surjectors.surjectors.affine_coupling_funnel import Funnel
+from surjectors.surjectors.lu_linear import LULinear
+
+
+class MLP(Funnel, hk.Module):
+    def __init__(self, n_keep, decoder, dtype=jnp.float32):
+        self._r = LULinear(n_keep, dtype, with_bias=False)
+        self._w_prime = hk.Linear(n_keep, with_bias=True)
+
+        self._decoder = decoder
+        super().__init__(n_keep, decoder)
+
+    def inverse_and_likelihood_contribution(self, y):
+        y_plus, y_minus = self.split_input(y)
+        z, jac_det = self._r.inverse_and_likelihood_contribution(y_plus)
+        z += self._w_prime(y_minus)
+        lp = self._decode(z).log_prob(y_minus)
+        return z, lp + jac_det
+
+    def _decode(self, array):
+        mu, log_scale = self._decoder(array)
+        distr = distrax.MultivariateNormalDiag(mu, jnp.exp(log_scale))
+        return distr
+
+    def forward_and_likelihood_contribution(self, z):
+        pass
diff --git a/surjectors/surjectors/rq_coupling_funnel.py b/surjectors/surjectors/rq_coupling_funnel.py
diff --git a/surjectors/surjectors/slice.py b/surjectors/surjectors/slice.py
@@ -1,9 +1,18 @@
 from jax import numpy as jnp
 
-from surjectors.funnel import Funnel
+from surjectors.surjectors.affine_coupling_funnel import Funnel
 
 
 class Slice(Funnel):
-    def __init__(self, n_keep, kind="inference_surjection"):
-        # TODO: implement decoder and encoder
-        super().__init__(kind, decoder, encoder, n_keep)
+    def __init__(self, n_keep, decoder, encoder=None, kind="inference_surjection"):
+        super().__init__(n_keep, decoder, encoder, kind)
+
+    def inverse_and_likelihood_contribution(self, y):
+        z, y_minus = self.split_input(y)
+        lc = self.decoder.log_prob(y_minus, context=z)
+        return z, lc
+
+    def forward_and_likelihood_contribution(self, z):
+        y_minus = self.decoder.sample(context=z)
+        y = jnp.concatenate([z, y_minus], axis=-1)
+        return y
diff --git a/surjectors/surjectors/surjector.py b/surjectors/surjectors/surjector.py
@@ -1,10 +1,10 @@
 from abc import abstractmethod
 
 from jax import numpy as jnp
-from surjectors.surjectors.transform import Transform
+from surjectors.surjectors._transform import Transform
 
 
-_valid_kinds = ["inference_surjector", "generative_surjector", "bijector"]
+_valid_kinds = ["inference_surjector", "generative_surjector", "bijector", "surjector"]
 
 
 class Surjector(Transform):
@@ -49,4 +49,4 @@ def encoder(self):
 
     @property
     def dtype(self):
-        return self._dtype
+        return self._dtype