flax

f00cadc4 - updated extracting_intermediates in flax docs