flax

fix missing rng arg implementation in wmt example

#368

Merged