Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Comparing Treex with Equinox

See original GitHub issue

I think it is natural to compare Treex with Equinox, as both are PyTree-based libraries. The README currently says

Other Pytree-based approaches like Parallax and Equinox do not have a total state management solution to handle complex states as encountered in Flax. Treex has the Filter and Update API, which is very expressive and can effectively handle systems with a complex state.

I assume the total state management solution refers to the Kind system in Treeo. However, the recent RFC indicates that we cannot use that with higher-level frameworks like Elegy. Suppose I want to use Elegy for training loop automation, is there any reason I should prefer Treex over Equinox?

Issue Analytics

  • State:closed
  • Created a year ago
  • Comments:5 (2 by maintainers)

github_iconTop GitHub Comments

cgarciaecommented, May 18, 2022

Stateful layers like BatchNorm are experimental in Equinox, but I think this itself is not a problem. The real issue is that Equinox is designed to work with arbitrary pytrees so its filters can really only take into account leaf values, this means you don’t have access any metadata including the field’s name, which in turn means you probably cant express something like “select all the batch stats and cache leaves”.

Treeo filters on the other hand operate over FieldInfo which have access to the following information:

class FieldInfo:
    name: Optional[str]
    value: Any
    kind: Type[Any]
    module: Optional[Tree]

So if leaf is part of a treeo.Tree the filter has access to field’s name, kind, and the actual module instance (tree would’ve been a better name for this), and of course the leaf value itself. This makes treeo/treex filters a bit more powerful.

cgarciaecommented, May 18, 2022

Its a very good question. I’ve been following Equinox closely as Treex was inspired by it, and talk every one in a while with Patrick. Ideally we only have a single Pytree Module library but as it stands here is the situation:

  • Treex main contribution was to add Kinds to Pytrees that are able to emulate Flax collections. This is needed for the advanced lifted transformations as you see on Flax: nn.{scan, vmap, etc}, it also simplifies a lot of code since e.g. selecting all BatchStats become a trivial filter.
  • Problem is that Kinds are a non-standard mechanism in Pytrees which could be called “tagged leaves”. The problem is that since not all Pytree have this capability it become harder to interop with Pytrees from other libraries. For Elegy this meant that you had to create wrapper objects that correctly annotated the Kinds so Elegy could query them. Equinox has focused on staying within the current limitations of Pytree’s, I think is very valuable to explore how far you can get this way for a NN library.

If you don’t need BatchNorm or other potentially stateful layers (anything that needs a 'cache') then Equinox is a good choice, I highly encourage you to try it out, else Treex was designed to try deal with these cases (although the state management situation for pytrees is not perfect). I will definitely try to add Equinox support for Elegy.

Read more comments on GitHub >

github_iconTop Results From Across the Web

Yellow Birch in the Southern Appalachians
The first year of study began by collecting data at sites where a range of different size yellow birch trees occur. Equinox collected...
Read more >
What's the Difference Between a Solstice and an Equinox?
Monkeys and apes are both primates, which means they're both part of the human family tree. As distinguished relatives, we should probably be...
Read more >
Equinox Comparison | SunPower by Precis | California Solar ...
How Solar Works - Understanding The Differences. SunPower® Equinox™ is the only all-in-one, complete home solar system.
Read more >
UKC Forums - Equinox Joshua Tree - UKClimbing
Equinox Joshua Tree ... From anyones' experience, if Equinox was in the Peak or Wales say, how hard would it be? What would...
Read more >
Freezing parameters - Equinox - Patrick Kidger
In this example, we demonstrate how the filtering of equinox.filter_value_and_grad can be customised -- in this case, to only train some parameters and...
Read more >

github_iconTop Related Medium Post

No results found

github_iconTop Related StackOverflow Question

No results found

github_iconTroubleshoot Live Code

Lightrun enables developers to add logs, metrics and snapshots to live code - no restarts or redeploys required.
Start Free

github_iconTop Related Reddit Thread

No results found

github_iconTop Related Hackernoon Post

No results found

github_iconTop Related Tweet

No results found

github_iconTop Related Post

No results found

github_iconTop Related Hashnode Post

No results found