Docs and switch to jr.key

[danielward27]

Changes:

• Now support (and enforce) new style jax keys, i.e. update any jr.PRNGKey(seed) calls with jr.key(seed).
• Improved documentation.
• LogNormal now implemented as a transformed Normal distribution, rather than a transformed StandardNormal (neater structure, and nicer to do log_normal.base_dist.loc, to see the location parameter, rather than going into the Chain bijection previously used.
• Docs now publish on pushes to main.