tl;dr samplers should be swappable and composable. for that we need a common interface

There's a lot of hot stuff in the pipeline re. MCTS, the vanilla sampler, etc.

One thing I'm afraid of is that there's going to be a lot of spaghetti involving bespoke/subtly different ways to call different samplers, which will make benchmarking and comparison painful.

I want to get ahead of this issue by defining a common interface to samplers. Since this is Python I think this should be a Protocol, something like:

class Sampler[Cfg, ST](Protocol):
     def __call__(gen_tokens: jax.Array, logits: jax.Array, attention_scores: jax.Array,
                  cfg: Cfg, state: ST) -> tuple[jax.Array, ST]:
          ...

which is a light touch (no need to inherit) but can still be checked. This should be a generic enough framework for people to be able to implement their favorite MuZero etc. and have it all plug in to the same harnesses.

My goal here is to have an easy to maintain sampler benchmarking suite with easy plug and play samplers.

EDIT

given the jax idiom of passing and returning state as an argument (and to support some sampler work of my own, tee hee), I think it will make sense to expand this interface to include a ST type var.

EDIT 2

since we're returning a jax.Array in place of the token output, I propose that it be acceptable to return either a single or an entire sequence of tokens at once from a sample call. the callers should be able to handle either case (and, really, a single token is just a sequence of length 1)