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.