flax

c117c29d - Cloned from CL 333372978 by 'g4 patch'.

Cloned from CL 333372978 by 'g4 patch'. Original change by qiuminxu@qiuminxu:flax_examples_xprof:9682:citc on 2020/09/23 14:26:02. test xprof on flax linen and profile scope PiperOrigin-RevId: 337964891