flax

b1ca884a - Handle passing Variables and Module instances as arguments in lifted transforms

Handle passing Variables and Module instances as arguments in lifted transforms This allows users to pass mutable variables and module instances into transformed methods, for instance for using local helper methods.