cgarciae / treeo

A small library for creating and manipulating custom JAX Pytree classes

https://cgarciae.github.io/treeo

MIT License

58 stars 4 forks source link

Closed cgarciae closed 3 years ago

cgarciae commented 3 years ago

Changes

Removes opaque_is_equal, same functionality available through opaque.
Adds compact decorator that enable the definition of Tree subnodes at runtime.
Adds the Compact mixin that adds the first_run property and the get_field method.