unflaxify#
Recursively transforms flax.linen
modules into Penzai layers.
This utility is designed to make it easier to:
Patch and inspect
flax.linen
Modules using Penzai tools.Simplify the process of migrating from
flax.linen
to Penzai by swapping out components one by one.
It relies on inspecting Flax internals and may not be completely robust. Transformed layers in particular are not currently supported.
Classes
Representation of the arguments and kwargs of a function call. |
|
A StateVariable label for a Flax variable. |
|
A representation of an intercepted Flax module method call. |
|
A frozen representation of data in a particular Flax scope. |
Functions
|
Creates an |