[Flax] Correct all return tensors to numpy
# What does this PR do?

This PR adapts all examples to return `numpy` instead of `jax` to avoid preventing asynchronous dispatch: https://jax.readthedocs.io/en/latest/async_dispatch.html