[docs] Functional API example code does not use Functional API functions

The docs [here](https://flax.readthedocs.io/en/latest/nnx_basics.html#the-flax-functional-api) state "The Flax NNX Functional API consists of three basic methods: [nnx.split](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), [nnx.merge](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge), and [nnx.update](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.update). ... Below is an example of of StatefulLinear [nnx.Module](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) that uses the Functional API."

But the code that follows does not use `nnx.split`, `nnx.merge`, or `nnx.update`.