flax

Make README and Module docstring examples runnable

#1328

Merged