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

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

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)

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?

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 😬