flax

d03a2622 - Added missing @ causing functools.wrap to have no effect in dynamic_scale.py