Enable shared mcmc parameters with tempered smc

[andrewdipper]

Allows for mcmc_parameters to be passed to the mcmc kernel as shared parameters prior to applying vmap. Thus shared parameters will not need to be duplicated for each individual particle.

This change filters mcmc_parameters by the length of the first dimension. Any parameters with length 1 are considered shared (note this is also acceptable in the case of just a single particle) and the rest are unshared. Shared parameters are then closed over before applying vmap so they don't need duplication. The behavior remains the same for shared parameters that are duplicated as they are just treated as unshared as before. This seems like the most reasonable way to handle shared parameters but let me know.

Related to #690
cc @ciguaran

Some notes:

• the some of tests in test_kernel_compatibility.py can use 1 instead of self.n_particles though the correct choice arguably depends on the meaning of the parameter - I changed the ones in test_inner_kernel_tuning.py that were apparently shared
• https://blackjax-devs.github.io/sampling-book/algorithms/TemperedSMC.html and https://blackjax-devs.github.io/sampling-book/algorithms/TemperedSMCWithOptimizedInnerKernel.html should probably change to avoid duplication of the shared parameters
• I couldn't find a way to filter shared/unshared parameters with jax.tree_util functions - None always remains in the structure but maybe I missed something

[junpenglao]

Please put them in a single for loop.

[andrewdipper]

Done

[junpenglao]

LGTM overall - is it still draft?

[andrewdipper]

The draft was mostly for if the tests should be modified - removed it

[ciguaran]

calling extend params shouldn't be needed anymore

[andrewdipper]

I'm still using extend_params since it needs to convert to an array and add a leading dimension of length one. I could have added a separate helper but it'd do the same thing as extend_params but just without the jnp.repeat part

[junpenglao]

I think it is worth to modify extend_params into something simple without the jnp.repeat.

Would this work?

def extend_params(params: ArrayLikeTree):
    return jax.tree.map(lambda a: jnp.expand_dims(a, 0), params)

[andrewdipper]

That would work, perhaps with potentially a jnp.asarray() included. The concern I have with modifying extend_params is I'd assume there's still a use case for it with initializing unshared parameters?

A current use case is for the remaining tests that are using duplicated parameters for testing the unshared case.

[junpenglao]

@ciguaran Is there any other use case besides dimension matching initially?

[ciguaran]

There's none.

[ciguaran]

I'd assume there's still a use case for it with initializing unshared parameters -> that is a good concern although I think that in most cases we wouldn't initialize by "extending" with copies. We would do it by sampling from some probability distribution. So you can modify it as Junpeng suggested.

[andrewdipper]

ah, that makes sense. thanks

[ciguaran]

I'd add a comment highlighting that here you are testing that extending with a copy is the same as having a 1 dimension parameter. Otherwise it may be tricky to get why this is even hapening.

[andrewdipper]

done

[junpenglao]

Thank you!