Fix order of state and input in Flax Quickstart README
# What does this PR do?
Correct the order of state and input for `flax.linen.apply()` in Flax Quickstart README

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

## Who can review?
@patrickvonplaten @patil-suraj 
Note: if the above correction is valid, then we ll need to update [this image](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/lm_flax_inference.png) as well 