[Fix doc example] fix missing import jnp
# What does this PR do?

Some doc examples in Flax models miss `import jax.numpy as jnp` (including the flax template).
This PR fixes it.
_The documentation is not available anymore as the PR was closed or merged._