Example ideas

I might take a stab at these myself; but here are some ideas for examples that might showcase the strengths of ninjax

- batch normalization; known to feel tricky to newcomers to jax
- activation norms or regularization terms that do not simplify to a plain weight decay

There are probably  others like it, but I thought id put these thoughts out here