Upgrade Flax NNX Gemma Sampling Inference doc

[8bitmp3]

Preview: https://flax--4325.org.readthedocs.build/en/4325/guides/gemma.html

Also fixes broken code after:

! git clone https://github.com/google/flax.git flax_examples

...
- sys.path.append("./flax_examples/flax/nnx/examples/gemma")
+ sys.path.append("./flax_examples/examples/gemma")
...

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[8bitmp3]

Sampler configuration (https://github.com/google/flax/blob/main/examples/gemma/sampler.py):

    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)

Throws an error

TypeError                                 Traceback (most recent call last)
...
in <cell line: 1>()
----> 1 sampler = sampler_lib.Sampler(
      2     transformer=transformer,
      3     vocab=vocab,
      4     params=params['transformer'],
      5 )

TypeError: Sampler.__init__() got an unexpected keyword argument 'params'

@cgarciae PTAL

[cgarciae]

Hey @8bitmp3! I cleaned up this guide a little bit. Can you take a look at the new version?

[8bitmp3]

thanks @cgarciae 👍
on it

[8bitmp3]

Reopening after #4334 fixes