nnx.Pytree - is this still been worked on?

[adam-hartshorne]

Am just wondering if this feature is still been worked on, it seems absolutely crucial if you already have a decent code base which uses for instance equinox Module's.

[nnx] add Pytree #4154

[cgarciae]

Hey @adam-hartshorne! I haven't forgotten about this. Curious what is your use case though?

[adam-hartshorne]

I already have all my codebase using equinox. However, I would like to implement the algorithm detailed in Appendix C.2 of the following paper.

https://arxiv.org/pdf/2106.03885

There is no obvious way to achieve this in equinox, i.e. store B variable that you keep updating, as you can't have a variable in a class which you update using self.B = new_B_value type operation.

As far as can see in nnx you can do this by default. However, such a class is not pure PyTree so I can't drop into my existing codebase easily.

The PR I linked appears to do exactly this with the nnx.Pytree wrapper.

[cgarciae]

I see. Note that if you embed an nnx.Pytree in an equinox Module you would still be bound by value semantics, meaning you have to manually propagate the new state updates (you must output the Pytree).