Static Typing

I'm planning to add as many static types as possible to the code, e.g. `sample(self, num_steps : int,...)`. This has two benefits:

1. Prevents many simple bugs
2. Makes the code much easier to read/use. For example, simply by looking at the type judgment `x_initial : str | jaxarray` , one would see that x_initial must be *either* a str or a jax array. 