Fix doctests (1/x)

Author: dirmeier

.gitattributes

@@ -0,0 +1 @@
+*.ipynb linguist-vendored

.pre-commit-config.yaml

@@ -20,15 +20,15 @@ repos:
     args: ["--ignore-missing-imports"]
     files: "(surjectors|examples)"
 
-- repo: https://github.com/jorisroovers/gitlint
-  rev: v0.19.1
-  hooks:
-  - id: gitlint
-  - id: gitlint-ci
-
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.3.0
   hooks:
     - id: ruff
       args: [ --fix ]
     - id: ruff-format
+
+- repo: https://github.com/jorisroovers/gitlint
+  rev: v0.18.0
+  hooks:
+    - id: gitlint
+    - id: gitlint-ci

.python-version

@@ -0,0 +1 @@
+3.11.9

README.md

@@ -82,12 +82,19 @@ Contributions in the form of pull requests are more than welcome. A good way to
 
 In order to contribute:
 
-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,
-5) submit a PR 🙂
-
+1) Clone `surjectors` and install  `uv` from [here](https://github.com/astral-sh/uv).
+2) Create a new branch locally `git checkout -b feature/my-new-feature` or `git checkout -b issue/fixes-bug`.
+3) Install all dependencies via `uv sync --all-groups`.
+4) Activate the virtual environment: `source .venv/bin/activate`.
+5) Install `pre-commit` and `gitlint` via:
+
+   ```shell
+   pre-commit install
+   gitlint install-hook
+   ``` 
+6) Implement your contribution and ideally a test case.
+7) Test it by calling `make format`, `make lints` and `make tests` on the (Unix) command line.
+8) Submit a PR 🙂.
 
 ## Citing Surjectors
 
@@ -109,4 +116,4 @@ If you find our work relevant to your research, please consider citing:
 
 ## Author
 
-Simon Dirmeier <a href="mailto:sfyrbnd @ pm me">sfyrbnd @ pm me</a>
+Simon Dirmeier <a href="mailto:simd23 @ pm me">simd23 @ pm me</a>

examples/autoregressive_inference_surjection.py

@@ -11,112 +11,112 @@
 from matplotlib import pyplot as plt
 
 from surjectors import (
-    AffineMaskedAutoregressiveInferenceFunnel,
-    Chain,
-    MaskedAutoregressive,
-    TransformedDistribution,
+  AffineMaskedAutoregressiveInferenceFunnel,
+  Chain,
+  MaskedAutoregressive,
+  TransformedDistribution,
 )
 from surjectors.nn import MADE, make_mlp
 from surjectors.util import as_batch_iterator, unstack
 
 
 def _decoder_fn(n_dim):
-    decoder_net = make_mlp([4, 4, n_dim * 2])
+  decoder_net = make_mlp([4, 4, n_dim * 2])
 
-    def _fn(z):
-        params = decoder_net(z)
-        mu, log_scale = jnp.split(params, 2, -1)
-        return distrax.Independent(distrax.Normal(mu, jnp.exp(log_scale)))
+  def _fn(z):
+    params = decoder_net(z)
+    mu, log_scale = jnp.split(params, 2, -1)
+    return distrax.Independent(distrax.Normal(mu, jnp.exp(log_scale)))
 
-    return _fn
+  return _fn
 
 
 def _made_bijector_fn(params):
-    means, log_scales = unstack(params, -1)
-    return surjectors.Inverse(surjectors.ScalarAffine(means, jnp.exp(log_scales)))
+  means, log_scales = unstack(params, -1)
+  return surjectors.Inverse(surjectors.ScalarAffine(means, jnp.exp(log_scales)))
 
 
 def make_model(n_dimensions):
-    def _flow(**kwargs):
-        n_dim = n_dimensions
-        layers = []
-        for i in range(3):
-            if i != 1:
-                layer = AffineMaskedAutoregressiveInferenceFunnel(
-                    n_keep=int(n_dim / 2),
-                    decoder=_decoder_fn(int(n_dim / 2)),
-                    conditioner=MADE(int(n_dim / 2), [8, 8], 2),
-                )
-                n_dim = int(n_dim / 2)
-            else:
-                layer = MaskedAutoregressive(
-                    conditioner=MADE(n_dim, [8, 8], 2),
-                    bijector_fn=_made_bijector_fn,
-                )
-            layers.append(layer)
-            # TODO(simon): needs to change order
-            # layers.append(Permutation(order, 1))
-        chain = Chain(layers)
-
-        base_distribution = distrax.Independent(
-            distrax.Normal(jnp.zeros(n_dim), jnp.ones(n_dim)),
-            reinterpreted_batch_ndims=1,
+  def _flow(**kwargs):
+    n_dim = n_dimensions
+    layers = []
+    for i in range(3):
+      if i != 1:
+        layer = AffineMaskedAutoregressiveInferenceFunnel(
+          n_keep=int(n_dim / 2),
+          decoder=_decoder_fn(int(n_dim / 2)),
+          conditioner=MADE(int(n_dim / 2), [8, 8], 2),
         )
-        td = TransformedDistribution(base_distribution, chain)
-        return td.log_prob(**kwargs)
+        n_dim = int(n_dim / 2)
+      else:
+        layer = MaskedAutoregressive(
+          conditioner=MADE(n_dim, [8, 8], 2),
+          bijector_fn=_made_bijector_fn,
+        )
+      layers.append(layer)
+      # TODO(simon): needs to change order
+      # layers.append(Permutation(order, 1))
+    chain = Chain(layers)
+
+    base_distribution = distrax.Independent(
+      distrax.Normal(jnp.zeros(n_dim), jnp.ones(n_dim)),
+      reinterpreted_batch_ndims=1,
+    )
+    td = TransformedDistribution(base_distribution, chain)
+    return td.log_prob(**kwargs)
 
-    td = hk.transform(_flow)
-    td = hk.without_apply_rng(td)
-    return td
+  td = hk.transform(_flow)
+  td = hk.without_apply_rng(td)
+  return td
 
 
 def train(rng_seq, data, model, max_n_iter=1000):
-    train_iter = as_batch_iterator(next(rng_seq), data, 100, True)
-    params = model.init(next(rng_seq), **train_iter(0))
+  train_iter = as_batch_iterator(next(rng_seq), data, 100, True)
+  params = model.init(next(rng_seq), **train_iter(0))
 
-    optimizer = optax.adam(1e-4)
-    state = optimizer.init(params)
+  optimizer = optax.adam(1e-4)
+  state = optimizer.init(params)
 
-    @jax.jit
-    def step(params, state, **batch):
-        def loss_fn(params):
-            lp = model.apply(params, **batch)
-            return -jnp.sum(lp)
+  @jax.jit
+  def step(params, state, **batch):
+    def loss_fn(params):
+      lp = model.apply(params, **batch)
+      return -jnp.sum(lp)
 
-        loss, grads = jax.value_and_grad(loss_fn)(params)
-        updates, new_state = optimizer.update(grads, state, params)
-        new_params = optax.apply_updates(params, updates)
-        return loss, new_params, new_state
+    loss, grads = jax.value_and_grad(loss_fn)(params)
+    updates, new_state = optimizer.update(grads, state, params)
+    new_params = optax.apply_updates(params, updates)
+    return loss, new_params, new_state
 
-    losses = np.zeros(max_n_iter)
-    for i in range(max_n_iter):
-        train_loss = 0.0
-        for j in range(train_iter.num_batches):
-            batch = train_iter(j)
-            batch_loss, params, state = step(params, state, **batch)
-            train_loss += batch_loss
-        losses[i] = train_loss
+  losses = np.zeros(max_n_iter)
+  for i in range(max_n_iter):
+    train_loss = 0.0
+    for j in range(train_iter.num_batches):
+      batch = train_iter(j)
+      batch_loss, params, state = step(params, state, **batch)
+      train_loss += batch_loss
+    losses[i] = train_loss
 
-    return params, losses
+  return params, losses
 
 
 def run(n_iter):
-    n, p = 1000, 20
-    rng_seq = hk.PRNGSequence(2)
-    y = jr.normal(next(rng_seq), shape=(n, p))
-    data = namedtuple("named_dataset", "y")(y)
+  n, p = 1000, 20
+  rng_seq = hk.PRNGSequence(2)
+  y = jr.normal(next(rng_seq), shape=(n, p))
+  data = namedtuple("named_dataset", "y")(y)
 
-    model = make_model(p)
-    params, losses = train(rng_seq, data, model, n_iter)
-    plt.plot(losses)
-    plt.show()
+  model = make_model(p)
+  params, losses = train(rng_seq, data, model, n_iter)
+  plt.plot(losses)
+  plt.show()
 
-    y = jr.normal(next(rng_seq), shape=(10, p))
-    print(model.apply(params, **{"y": y}))
+  y = jr.normal(next(rng_seq), shape=(10, p))
+  print(model.apply(params, **{"y": y}))
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--n-iter", type=int, default=1_000)
-    args = parser.parse_args()
-    run(args.n_iter)
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--n-iter", type=int, default=1_000)
+  args = parser.parse_args()
+  run(args.n_iter)