Link with Equinox

[pofmind]

As of now, it seems that this project overlaps with Equinox on many ideas.

Hence the question, why to not have worked on Equinox directly instead ?

[danieldjohnson]

Thanks for the question!

In terms of tooling, both Equinox and Penzai are pretty modular, so Penzai's utilities should work out of the box with Equinox models and vice versa. This includes the treescope pretty-printer, the selector system, and the named axis library.

For the neural network components, Penzai has an opinionated set of conventions that differ from those in Equinox. In particular, penzai.nn is more restrictive about how layers are composed, and handles parameters differently (using sentinel node types instead of filtered transformations). These changes are designed to expose as much as possible in the pytree structure of the model, going beyond what conventional Equinox models allow. This means it wouldn't be straightforward to directly extend Equinox, since this would require changing some of Equinox's core design ideas.

(Early in development, Penzai layers actually were Equinox modules under the hood, but we didn't use any of Equinox's neural network components or transformations so ultimately we separated out the base class as well. This also allows us to support a bit more customization in Penzai's explicit dataclass decorator relative to Equinox's implicit transform. Penzai structs/layers and Equinox modules are both just pytree dataclasses, so they can be used in very similar ways.)

The goal of Penzai is to be a toolkit that you can use as much or as little of it as you want, without it getting in your way. If you don't want to follow the conventions of penzai.nn, you don't have to.