flax

ea47d81c - Renames `train_state` to `state`.

Renames `train_state` to `state`. As @8bitmp3 has mentioned in [a comment](https://github.com/google/flax/commit/c4b3464427ae1d7e12ef3e114dd8ed4022709ddf#commitcomment-50892606) currently it's hard to find documentation about the new `state.apply_gradients()` because we call the variable `train_state` instead of `state` as in the examples. Note that we decided to go for simply `state` in the examples because we stored the new `TrainState` in a module called `flax.training.train_state`.