Clarification on how particles (eg. in SMC and SVGD) should be handled.

[antotocar34]

In smc/base.py, the particles object is typed to be a PyTree.

However, the following line assumes a more specific structure

blackjax/blackjax/smc/base.py

Line 111 in 25eee35

num_particles = jax.tree_util.tree_flatten(particles)[0][0].shape[0]

Specifically, the above line assume that particles has to have the following structure (let's say you're sampling a univariate distribution):

[ [1.], [1.2] ]

Unless I'm missing something, this does not generalize to particles that look like:

[ {"theta": [1.]}, {"theta": [1.2]} ]

Perhaps a workaround would be to type particles as List[PyTree]?
The following functions could convert the particles to vmap-able arrays with the following functions

from jax.flatten_util import ravel_pytree

def ravel_particles(particles: List[PyTree]) -> Tuple[Array, Callable]:
    _, unravel_fn = ravel_pytree(particles[0])
    return jnp.array([ravel_pytree(p)[0] for p in particles]), unravel_fn

def unravel_particles(particles: Array, unravel_fn: Callable) -> List[PyTree]:
    return jax.vmap(unravel_fn)(particles)

[rlouf]

I'm not in front of my computer right now, what's the output of flatten_pytree on both example data structures?

[antotocar34]

Oops, they both return

[1.0, 1.2]

I was wrong when I said that the num_particles line doesn't work on the dictionary.
Nevermind that :) . There should however be the bigger issues that the num_particles line
fails on a PyTree, and only works if particles is an Array.

If I pass in [ {"theta": [1.]}, {"theta": [1.2]} ] to the one_step function, the

blackjax/blackjax/smc/base.py

Line 111 in 25eee35

num_particles = jax.tree_util.tree_flatten(particles)[0][0].shape[0]

line will raise an error.

To see why, note the following:

[ins] In [1]: import jax.numpy as jnp

[ins] In [2]: from jax.tree_util import tree_flatten

[ins] In [3]: x = [ [1., 4.], [1.2, 3.] ]

[ins] In [4]: y = jnp.array(x)

[ins] In [5]: tree_flatten(x)
Out[5]: ([1.0, 4.0, 1.2, 3.0], PyTreeDef([[*, *], [*, *]]))

[ins] In [6]: tree_flatten(y)
Out[6]: 
([Array([[1. , 4. ],
         [1.2, 3. ]], dtype=float32)],
 PyTreeDef(*))

Notice how tree_flatten returns a jax array in a list.

Running the following code (copied from examples documentation, changed initial particles to be list of dictionaries):
https://gist.github.com/antotocar34/ebfc1a8a53ca17af2383c4af2f0c2e52

yields the following error:

File ~/Documents/projects/blackjax/blackjax/smc/ess.py:82, in ess_solver(logdensity_fn, pa
rticles, target_ess, max_delta, root_solver, use_log_ess)
     49 def ess_solver(
     50     logdensity_fn: Callable,
     51     particles: PyTree,
   (...)
     55     use_log_ess: bool = True,
     56 ):
     57     """Build a Tempered SMC step.
     58 
     59     Parameters
   (...)
     80 
     81     """
---> 82     n_particles = jax.tree_util.tree_flatten(particles)[0][0].shape[0]
     84     logdensity = logdensity_fn(particles)
     85     if use_log_ess:

IndexError: tuple index out of range

[rlouf]

So there is a funny thing with JAX that took me a while to get used to. In this case you would have to pass:

import jax.numpy as jnp

particles = {'weight': jnp.array([1.2, 2.])}

to represent two particles. This is because of the way vmap works on PyTrees. Try the following:

import jax

def my_fun(particle):
   weight = particle['weight']
   return weight + 1.

res = jax.vmap(my_fn)(particles)
print(res)

[antotocar34]

Ah ok, that makes sense. So then particles must be a PyTree, with the leaves of that tree being an array of size num_particles.
Good to know thanks a million!

Where do you think I could add a section about that in the documentation?

[rlouf]

Where do you think I could add a section about that in the documentation?

The best place is probably in the definition of SMCState which I added in #279, and in the definition of SVGDState.

...then we'll need to include these in the documentation 😬

[ciguaran]

I've added docs in SMCState