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

ArgsAndKwargs	Representation of the arguments and kwargs of a function call.
InterceptedFlaxModuleMethod	A representation of an intercepted Flax module method call.
InterceptedFlaxScopeData	A frozen representation of data in a particular Flax scope.

Functions

unflaxify_apply(module, variables, *dummy_args)

Creates an InterceptedFlaxModuleMethod from applying a Flax module.