flax

Use '_rng' for init rng arg to allow apply to have an 'rng' arg.

#155

Merged