Jaxns Roadmap

[Joshuaalbert]

UPDATED

9 Jan, 2024

2.4.0 is now out. JAXNS has evolved a lot.

1. The current mission is to try to get more groups across the sciences using it for their analysis. To that end I want to make it easier to use, more flexible, faster, and more powerful for those really hard problems. To this end the underlying probabilistic programming framework has undergone lots of upgrades and been made stand alone. I.e. you can use it for numerous thinks besides nested sampling.
2. Phantom-powered nested sampling is now integrated. While incorporating phantom points doesn't speed up evidence calculations, it does provide exponentially faster parameter estimations. So posterior estimation is significantly improved.
3. In addition, I've been paving the way to merge AI+deep learning with JAXNS, e.g. to allow learning priors, improving efficiency of sampling, and more. I expect I'll make a paper on this in the coming months to showcase this powerful new direction, and how it can assist in science.

2 June, 2022

Release 1.1.0 is out. Asymptotic accuracy, everything has parallelised option, everything used adaptive nested sampling, dynamic sampler works out of the box, fewer hyper parameters, works for arbitrary non-convex and discrete likelihoods and plateaus.

Next release 1.2 will tackle non-JAX substrates to enable a broader class of programs to take advantage of jaxns.

30 May, 2022

Release 1.0.1 is out, and once a few minor improvements are made will move to 1.1.0.

5 May, 2022

Release plans for 1.1 and 1.2 are available. They will cover increased likelihood efficiency and more applicability for a wide range of projects.

3 March, 2022

JAXNS major release 1 is now officially released with 1.0.0.

The basic plan is that I'll be listening to users and making small improvements with a minor release every month or so.

What can you do?

• If you find a bug, check out the issues to see if it's reported already, or make a new one.
• If you would like a new feature make a feature request.

[Joshuaalbert]

Release Candidate 1.0 is out!

[jecampagne]

One of the first challenge is to setup a proper icdf function needed by NS...

[Joshuaalbert]

The support of x is R^2? or would you like it compact like [-a,a] x [-a,a]? In any case, are you looking to treat the dual moon as a prior?

This example, shows how to create a prior when you only know the log-probability and support.

https://github.com/Joshuaalbert/jaxns/blob/master/examples/simple/priors_without_quantiles.ipynb

(Note: for some reason when I save ipynb files for the examples, they don't commit properly and github shows cells and old code sometimes. Still trying to figure out that, without modifying raw ipynb code)

I'd look to incorperate dual moon, as it looks indeed nice.

[jecampagne]

Well looking at THIS I was trying to mimic the posterior plot... May be it is not ad equat. At least if one want to get samples from this distribution, HMC is not doing very well.

[Joshuaalbert]

JAXNS does quite alright without any tuning. Thanks for sharing the model. It actually revealed a bug, which I think is an edge case #52 (need to take a deeper look).

Dual Moon example
https://github.com/Joshuaalbert/jaxns/blob/master/examples/toy/dual_moon.ipynb

[jecampagne]

Thanks that sounds good

[Joshuaalbert]

1.0.0 is released!

[Joshuaalbert]

@fehiepsi JAXNS 1.1.0 is out. It's quite robust now and should fulfill many of numpyro's users' difficult usecases. We also have a physics project that has taken it up and will be vetting it against some truly difficult cosmological problems.

[fehiepsi]

@Joshuaalbert I tried to bump jaxns to 1.1.0 here but it seems to make the gaussian_shells example hang:

CI / examples (3.7) (pull_request) Cancelled after 360m — examples (3.7)

Could you take a look?

[Joshuaalbert]

Hi @fehiepsi I would love to check it out. It's strange because that problem is easily silenced with jaxns. Can you point me to where in your code you call into jaxns.NestedSampler?

[fehiepsi]

@Joshuaalbert We made a wrapper, which called jaxns.NestedSampler at this line.

[Joshuaalbert]

Re here we don't support marginalised keywork in the sampling call anymore. The user can import the utility script to do that and run on the generated samples after.

Re here I'm not sure this is needed anymore because at this line we will prepare the likelihood definition so that we discard any prior parameters that are not needed. So if your likelihood needs x and y, but you prior also contains hierarchical priors a and b then you can define simply def log_likelihood(x,y) and we will automatically discard the a and b at call time.

Re here you may prefer to use the construct,

with PriorChain() as prior_chain:
  for name in param_names:
    prior = UniformPrior(name + "_base", prototype_trace[name]["fn"].shape())
    prior_chain.push(prior)

static JAX code violations

Re here Selecting this slice is non-static which means not jit. Mind you, the sampler runs pretty fast even without jit because it makes use of scan and while_loop which naturally jit there subcomponents. You just won't get the whole think compiling top-down.

Re here I may be missing something, but I don't see you calling jaxns.resample(...) first to produce equally weighted samples? Note, resample(...) does require a static number of samples to produce, however there is a way to do it in a static way. If your Predictive class could be made to take into account weights, e.g. when doing marginalisation, then you wouldn't need to resample.

where some improvement could be made on jaxns side

Re here it looks like the place where a lot of difference can be made is in default construction and evaluation parameters. I think it would be best if I set these to be more robust, so that at least things always terminate in reasonable time, and give a simple few rules of thumb for increase accuracy depending on the user's goal. That's sort of what I'm doing now, but I can do better.

[Joshuaalbert]

Sorry this took so long to get to. I got married in July and am just now getting back :)

[fehiepsi]

don't support marginalised keywork in the sampling call anymore
Thanks, we will clean up that comment. Currently, we do not maintain signatures in the wrapper. Users need to use the signatures in jaxns.

Re signature of log-likelihood: here we don't know signatures of a likelihood in advance (it can be varied depending on the users' model definition) so we need that trick to prepare a likelihood with the desired signature for jaxns.

Re using context manager: thanks! just curious: is there any difference when using PriorChain() as a context manager?

Selecting this slice is non-static which means not jit.

You are right. Currently, I'm not sure if we have a usage case that requires jitting the run code. I'll keep it in mind when there's a user request for it.

calling jaxns.resample(...) first to produce equally weighted samples?

That predictive code is only used to transform samples back from unit cubes. Users will need to use methods get_samples or get_weighted_samples for predictions.

it would be best if I set these to be more robust

Thanks! Hopefully, the gaussian_shells example will be able to run with the new settings. Currently, it hangs for at least 6 hours making us unable to update jaxns to 1.1.0.

I got married in July and am just now getting back

Congratulation and welcome back!

[Joshuaalbert]

We should see 1.2 released at the end of September.

That will introduce other substrates to enable non-JAX computation. Initially, numpy and pure python backends will be released to open up computation to many physics and biology projects where simulations are written in some other language will cython bindings.

[Joshuaalbert]

Release 1.2 has been abandoned and a new 2.0 release is instead coming. There will be breaking API changes, for a good reason, and there will be better stability and performance promises. In this major release series we'll be focusing on:

1. ease of use
2. enabling more science domains
3. improving the likelihood constraint sampling efficiency

2.0 in prep now and can be tracked here: #71