Add GitHub actions

Add a GitHub action to compute code coverage

Author: dirmeier

.github/workflows/ci.yaml

@@ -77,3 +77,7 @@ jobs:
       - name: Run tests
         run: |
           hatch run test:test
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v3
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

README.md

@@ -2,6 +2,7 @@
 
 [![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/surjectors/actions/workflows/ci.yaml)
+[![codecov](https://codecov.io/gh/dirmeier/surjectors/branch/main/graph/badge.svg?token=dn1xNBSalZ)](https://codecov.io/gh/dirmeier/surjectors)
 [![version](https://img.shields.io/pypi/v/surjectors.svg?colorB=black&style=flat)](https://pypi.org/project/surjectors/)
 
 > Surjection layers for density estimation with normalizing flows
@@ -19,13 +20,34 @@ Surjectors makes use of
 - Optax for gradient-based optimization,
 - JAX for autodiff and XLA computation.
 
-## Documentation
+## Examples
 
-Documentation can be found [here](https://surjectors.readthedocs.io/en/latest/).
+You can, for instance, construct a simple normalizing flow like this:
+
+```python
+import distrax
+from jax import numpy as jnp
+from surjectors import Slice, LULinear, Chain
+from surjectors import TransformedDistribution
+from surjectors.nn import make_mlp
+
+def decoder_fn(n_dim):
+    def _fn(z):
+        params = make_mlp([32, 32, n_dim * 2])(z)
+        means, log_scales = jnp.split(params, 2, -1)
+        return distrax.Independent(distrax.Normal(means, jnp.exp(log_scales)))
+    return _fn
+
+base_distribution = distrax.Normal(jnp.zeros(5), jnp.ones(5))
+transform = Chain([Slice(10, decoder_fn(10)), LULinear(5)])
+pushforward = TransformedDistribution(base_distribution, transform)
+```
 
-## Examples
+More self-contained examples can be found in [examples](https://github.com/dirmeier/surjectors/tree/main/examples).
+
+## Documentation
 
-You can find several self-contained examples on how to use the algorithms in `examples`.
+Documentation can be found [here](https://surjectors.readthedocs.io/en/latest/).
 
 ## Installation
 
@@ -47,11 +69,11 @@ pip install git+https://github.com/dirmeier/surjectors@<RELEASE>
 ## Contributing
 
 Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled
-`"good first issue" <https://github.com/dirmeier/surjectors/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22>`_.
+[good first issue](https://github.com/dirmeier/surjectors/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22).
 
 In order to contribute:
 
-1) Clone Surjectors and install `hatch` via `pip install hatch`,
+1) Clone `Surjectors` and install `hatch` via `pip install hatch`,
 2) create a new branch locally `git checkout -b feature/my-new-feature` or `git checkout -b issue/fixes-bug`,
 3) implement your contribution and ideally a test case,
 4) test it by calling `hatch run test` on the (Unix) command line,

docs/index.rst

@@ -14,26 +14,27 @@ Surjectors makes use of
 - Optax for gradient-based optimization,
 - JAX for autodiff and XLA computation.
 
-Example usage
--------------
+Example
+-------
 
 You can, for instance, construct a simple normalizing flow like this:
 
     >>> import distrax
-    >>> from jax import random as jr, numpy as jnp
+    >>> from jax import numpy as jnp
     >>> from surjectors import Slice, LULinear, Chain
     >>> from surjectors import TransformedDistribution
+    >>> from surjectors.nn import make_mlp
     >>>
     >>> def decoder_fn(n_dim):
     >>>     def _fn(z):
-    >>>         params = make_mlp([4, 4, n_dim * 2])(z)
-    >>>         mu, log_scale = jnp.split(params, 2, -1)
-    >>>         return distrax.Independent(distrax.Normal(mu, jnp.exp(log_scale)))
+    >>>         params = make_mlp([32, 32, n_dim * 2])(z)
+    >>>         means, log_scales = jnp.split(params, 2, -1)
+    >>>         return distrax.Independent(distrax.Normal(means, jnp.exp(log_scales)))
     >>>     return _fn
     >>>
-    >>> base_distribution = distrax.Normal(jno.zeros(5), jnp.ones(1))
-    >>> flow = Chain([Slice(10, decoder_fn(10)), LULinear(5)])
-    >>> pushforward = TransformedDistribution(base_distribution, flow)
+    >>> base_distribution = distrax.Normal(jnp.zeros(5), jnp.ones(1))
+    >>> transform = Chain([Slice(10, decoder_fn(10)), LULinear(5)])
+    >>> pushforward = TransformedDistribution(base_distribution, transform)
 
 The flow is constructed using three objects: a base distribution, a transformation, and a transformed distribution.
 
@@ -64,7 +65,7 @@ Contributions in the form of pull requests are more than welcome. A good way to
 
 In order to contribute:
 
-1) Clone Surjectors and install :code:`hatch` via :code:`pip install hatch`,
+1) Clone :code:`Surjectors` and install :code:`hatch` via :code:`pip install hatch`,
 2) create a new branch locally :code:`git checkout -b feature/my-new-feature` or :code:`git checkout -b issue/fixes-bug`,
 3) implement your contribution and ideally a test case,
 4) test it by calling :code:`hatch run test` on the (Unix) command line,

docs/notebooks/introduction.ipynb

@@ -3,7 +3,11 @@
   {
    "cell_type": "markdown",
    "id": "4f6f4229-6b15-4e2b-89af-8957708479d7",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "# Constructing normalizing flows\n",
     "\n",
@@ -15,7 +19,10 @@
    "execution_count": 4,
    "id": "9497e202-3f4e-4602-9e90-669545f18816",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [],
    "source": [
@@ -28,7 +35,11 @@
   {
    "cell_type": "markdown",
    "id": "bfd96e3f-ef84-454f-a611-0cdc0a629d2d",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## How to construct a Haiku module\n",
     "\n",
@@ -40,7 +51,10 @@
    "execution_count": 30,
    "id": "c9db6163-6cb8-4968-8e3d-90d35a5094cf",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [],
    "source": [
@@ -53,7 +67,10 @@
    "execution_count": 61,
    "id": "f78c2fc0-f4f6-476d-8b15-12b4a874ce1d",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [],
    "source": [
@@ -88,7 +105,10 @@
    "cell_type": "markdown",
    "id": "b41e89a1-6df4-4069-a606-4f2412b9dc6a",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%% md\n"
+    }
    },
    "source": [
     "Constructing a Haiku module needs to be done within a `hk.transform` block. This can either be done by providing a function like here and an object. In our case we are using `hk.transform` on `pushforward(**kwargs)` which calls\n",
@@ -98,7 +118,11 @@
   {
    "cell_type": "markdown",
    "id": "689f867c-2259-4713-81b6-5352837cb342",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "We can now initialize the flow. Let's define a random data set first and then initialize the parameters."
    ]
@@ -108,7 +132,10 @@
    "execution_count": 65,
    "id": "cf91f2ce-6ba9-438c-95a7-ec9dcfbfea17",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [
     {
@@ -150,7 +177,11 @@
   {
    "cell_type": "markdown",
    "id": "4c8f8b8e-4e81-4b54-a1c4-72027db33e1c",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "The only trainable paramaters that are flow defines are the weights of the MLP. The MLP is used to compute the conditional probability density inside the `decoder_fn` function. \n",
     "The `Slice` surjector itself doesn't have paramters."
@@ -159,7 +190,11 @@
   {
    "cell_type": "markdown",
    "id": "84d8368a-2c75-4e72-a3b4-45f4cfc71fbf",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "We can now test the flow. Let's sample some data first."
    ]
@@ -169,7 +204,10 @@
    "execution_count": 71,
    "id": "77ae282d-9e30-49c9-98b1-6f9494b34a21",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [
     {
@@ -196,7 +234,11 @@
   {
    "cell_type": "markdown",
    "id": "04fd204d-f0c2-468b-a64d-ec292b071f34",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "As mentioned above, in order to dispatch to a method, we just provide a keyword argument. In this case this is `method='sample'`. Computing the log probability of the data can be done, by changing the method argument to `log_prob`."
    ]
@@ -206,7 +248,10 @@
    "execution_count": 72,
    "id": "2190c369-4242-45ec-9ac9-b9fa5bccc1dd",
    "metadata": {
-    "tags": []
+    "tags": [],
+    "pycharm": {
+     "name": "#%%\n"
+    }
    },
    "outputs": [
     {
@@ -227,7 +272,11 @@
   {
    "cell_type": "markdown",
    "id": "8c046ffd-a2b9-4348-b8f3-85e978b7998f",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## How to construct `TransformedDistribution` objects\n",
     "\n",
@@ -238,7 +287,11 @@
    "cell_type": "code",
    "execution_count": null,
    "id": "10c16d0c-0158-4845-b192-dd97967d7eb6",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": []
   }

docs/requirements.txt

@@ -7,6 +7,7 @@ seaborn
 sphinx
 sphinx-autobuild
 sphinx-book-theme
+sphinx-copybutton
 sphinx-math-dollar
 sphinx_autodoc_typehints
 sphinx_design

surjectors/__init__.py

@@ -2,7 +2,7 @@
 surjectors: Surjection layers for density estimation with normalizing flows
 """
 
-__version__ = "0.2.4"
+__version__ = "0.3.0"
 
 from surjectors._src.bijectors.lu_linear import LULinear
 from surjectors._src.bijectors.masked_autoregressive import MaskedAutoregressive