v17.0.2

Updates to the workflow packages led to some errors in the release to PyPi, which was fixed with this release. Below are the changes from v16-v17, which are of more interest to users.

Breaking Changes

• The AbstractBijection class now implements transform and inverse methods by indexing the outputs of transform_and_log_det and inverse_and_log_det. AFAIK, under JIT, dead code elimination prevents computation of log_det where unnecessary, minimizing the benefits of directly implementing these methods. Custom bijections should generally avoid overriding transform and inverse.
• The flowjax.wrappers module has been removed. Its functionality is now available in a new, separate package: paramax. Most functionality is unchanged when imported from paramax.
• Partial has been renamed to Indexed, as the former was likely to be confused with functools.partial or jax.tree_util.Partial.
• The deprecated fit_to_variational_target function has been removed. Use fit_to_key_based_loss instead.
• Numerical inverse methods are now provided using the NumericalInverse composition. This means:
  • Users directly calling inverse methods on BlockAutoregressiveNetwork will encounter an error. To resolve this, supply an
    inverse method via NumericalInverse (as is done in block_neural_autoregressive_flow).
  • Users accessing the attributes of block_neural_autoregressive_flows might also face a breaking changes; an additional .bijection may be required to extract BlockAutoregressiveNetwork from NumericalInverse.
• The uniform distribution is now non-trainable by default. Optimization of a uniform distribution seemed most commonly a mistake leading to e.g. violating support assumptions.
• The root-finding algorithms in flowjax.bisection_search have been moved to flowjax.root_finding. Check the updated function names and documentation in the new module if you directly use these methods.

---

Apologies for any inconvenience caused by these breaking changes. If you encounter issues or have questions, please feel free to open an issue.

---

What's Changed

• Update docs by @danielward27 in #190
• Paramax by @danielward27 in #192
• Provide transform and inverse method implementations on AbstractBijection by @danielward27 in #193
• Root finding by @danielward27 in #195
• Rename partial and rm fit to variational by @danielward27 in #196
• Root finding by @danielward27 in #198
• Make uniform non-trainable by @danielward27 in #199

Full Changelog: v16.0.0...v17.0.0