0.6.1

Changes

• Updates treeo to version 0.0.7

0.6.0

Shape Inference + @compact support 🎉

Changes

• Adds the tx.next_key() function and the tx.rng_key() context manager.
• Module.init new has the following behavior:
  • Accepts an optional inputs argument and runs the forward method if given.
  • Set the given key in the context so tx.next_key() can be used.
  • Accepts a call_method: str which defines the method to call, "__call__" used by default.
• Modules will now be initialized if constructed within @tx.compact functions when called by init.
• Adds @tx.compact_module decorator that can turn any function into a Module with a compact __call__ as the decorated function.
• New Crossentropy loss that generalizes BinaryCrossentropy, CategoricalCrossentropy and SparseCategoricalCrossentropy.
• New Flatten layers.

0.5.0

Major Changes

• Treex now depends on Treeo to generate its Pytree.
• update is now called merge consistent with Treeo, it also avoids name clashes with Optimizer.
• Kinds are no longer annotations, instead uses Tree's kind system. So this annotation

w: tx.Parameter[jnp.ndarray]

becomes

w: jnp.ndarray = tx.Parameter.node()

0.4.0

Changes

• Optimizer now flattens its params to be agnostic to the static components of the pytree.
• Generic types containing TreeParts are no longer valid type annotation as types like Tuple[int, tx.State[int]] make it appear as if the first element of the tuple where static and the second dynamic, when in fact Treex would treat to whole field as dynamic. Now your only option is tx.State[Tuple[int, int]].
• Adds the tx.Hashable class to wrap non-hashable types like numpy or jax arrays when you want to use them in static fields of a TreeObject.
• Adds FlaxModule: can wrap any Flax Module into a Treex Module.
• tabulate now accepts a sample_input and will show the input and output columns.
• Refactors a lot of the functional API.
• Introduces .freeze(), .unfreeze() and .frozen similar to train/eval/training.
• Updates BatchNorm and Dropout to leverage .frozen.

0.3.0

Changes

• TreeParts are now generic and statically behave like Union.
• filter now also accepts predicates.
• Optimizer.update was renamed to apply_updates.
• Expanded TreePart hierarchy.
• Added RngSeq Module (generates PRNGKeys on demand).
• TreeObject now has a metaclass that checks that super().__init__() is always called.
• Fields with TreeObject values are now automatically annotated if an annotation is not provided by the user.
• filter now also accepts predicates of the type FieldInfo -> bool.
• Added the Static annotation for when you want a field to be explicitly marked as a static part of the Pytree. This is useful if the field will hold a TreeObject but you don't want it to be a child of the Pytree e.g. ignored_linear: tx.Static[tx.Linear]

0.2.3

0.2.3

0.2.2

0.2.2

0.2.1

update version to 0.2.1