flax

Seq2Seq: Add back a bit of nn.sochastic now that `init_shape` issue is resolved.

#168

Merged