Joss Review: Example Usage

[animikhaich]

I installed Jax CPU, followed by Surjectors. A clean install inside a docker container.

Once installed, ran the example for constructing a simple normalizing flow as mentioned in the Readme#examples. I got the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/surjectors/bijectors/lu_linear.py", line 16, in __init__
    super().__init__(n_keep, None, None, "bijection", dtype)
  File "/usr/local/lib/python3.10/dist-packages/surjectors/surjectors/surjector.py", line 24, in __init__
    raise ValueError(
ValueError: 'kind' argument needs to be either of: inference_surjector/generative_surjector/bijector/surjector

It failed in the line: transform = Chain([Slice(10, decoder_fn(10)), LULinear(5)])

Link to the Review Thread: openjournals/joss-reviews#6188

[dirmeier]

Hi, my bad, I hadn't updated the pip version to v0.3.0. Assuming you are running version v0.3.0, the examples should run (at least the GitHub actions run and the documentation builds).

I updated the example in the README though to make it more friendly for users w/o Haiku background. Could you test this:

pip install surjectors==0.3.0

and then

import distrax
import haiku as hk
from jax import numpy as jnp, random as jr
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

@hk.without_apply_rng
@hk.transform
def flow(x):
    base_distribution = distrax.Independent(
        distrax.Normal(jnp.zeros(5), jnp.ones(5)), 1
    )
    transform = Chain([Slice(5, decoder_fn(5)), LULinear(5)])
    pushforward = TransformedDistribution(base_distribution, transform)
    return pushforward.log_prob(x)

x = jr.normal(jr.PRNGKey(1), (1, 10))
params = flow.init(jr.PRNGKey(2), x)
lp = flow.apply(params, x)
print(lp)

Thanks and cheers,
Simon

[dirmeier]

I think the example in the README is not very illustrative though. The ones in examples or here are better.

[animikhaich]

@dirmeier Thanks, that fixed it. But I found another minor bug in the documentation.

The Example in the Docs: https://surjectors.readthedocs.io/en/latest/#example is missing this import line: from surjectors import Slice, LULinear, Chain, which is causing the example to fail. Please update the documentation with the import.

Thanks!

[dirmeier]

Thanks! I made a PR #29 with the fix. Ill merge it after the review is done (I dont wanna add too many non-functional commits).

[animikhaich]

Understood Thanks! Closing this issue then and giving the go-ahead from my end. Good Luck! 👍🏻