Update wrappers

[danielward27]

This introduces a breaking change to reparameterizations.

Breaking changes:

• flowjax.wrappers.Lambda has been renamed to flowjax.wrappers.Parameterize
• flowjax.wrappers.BijectionReparam has been removed.
• flowjax.wrappers.Where has been removed.

Reasons for breaking change:

1. Lambda was renamed to Parameterize because Lambda is uninformative.
2. Parameterize is relatively simple, and flexible enough to cover the common uses of BijectionReparam and Where (see below).
3. Before, there was dependencies between the wrappers module, and the bijections, and vice verse (due to BijectionReparam) which made it challenging to avoid circular dependencies. The removal of BijectionReparam eliminates this issue.

Transitioning:

• Replace Lambda with Parameterize.
• Instead of e.g. BijectionReparam(scale, SoftPlus()), we can use Parameterize(softplus, inv_softplus(scale)), using jax.nn.softplus and flowjax.utils.inv_softplus.
• Instead of e.g. creating mask with Where(mask, weight, 0), we can use Parameterize(jnp.where, mask, weight, 0).

Apologies for any inconvenience and let me know if you have any specific use cases that this causes any issues for.