v0.5.9

• Use jax.tree_util.* functions to avoid warnings.
• Full Changelog: v0.5.8...v0.5.9

v0.5.8

• Fix a multithread bug due to using __slots__ attribute for threading.local objects.

v0.5.7

• Fix an EMA bug of handling boolean ndarray.

Full Changelog: v0.5.6...v0.5.7

v0.5.6

• Fix a batchnorm bug when create_offset=False.

Full Changelog: v0.5.5...v0.5.6

v0.5.5

Fix a bug in getting wrong level of abstraction due to jax recent updates.

Full Changelog: v0.5.4...v0.5.5

v0.5.4

• A new example on Adversarial Examples.
• Add feature_group_count optional argument for Conv modules.

Full Changelog: v0.5.3...v0.5.4

v0.5.3

• New VanillaRNN module.
• Improve WaveRNN example with dilated CNN.
• Initialize PAX module using __new__ method. This ensures dataclass modules are initialized correctly.

v0.5.2

• Deleting Pytree attributes are not allowed to prevent bugs.
• Expose EmptyNode to the public API. Users can use it to mark uninitialized state arrays.
• RngSeq module will refresh its internal key even in eval mode.
  Prevent bugs when a RngSeq module is used many times in eval mode.
• EMA module requires allow_int=True to allow integer leaves in the values.
  Integer leaves will be updated to the newest values instead of averaging.

v0.5.1

• Add state_dict and load_state_dict methods.
• Add pax.purecall function. Prefer this to the pax.module_and_value function.

v0.5.0

Two major changes in this release:

• Rewrite Module class:

  • We remove the name_to_kind dictionary and use a pytree_attributes list to manage the pytree.
  • Because there is no attribute's kind information, the user needs to implement the parameters method themself.
  • PAX provides utility function parameters_method that returns custom parameters methods.
  • Attributes can be a mix of ndarray/module leaves and other values. However, this will slow down (~3x) the tree_flatten and tree_unflatten methods. We do not recommend this practice.

• Remove .nn from public API:

  • Instead of pax.nn.Linear, it now is pax.Linear, etc.