Deprecate flax.jax_utils.replicate

`flax.jax_utils.replicate` should no longer be necessary and superceded by `jax.device_put_replicated`. We need to update our tests and use-cases accordingly, and add a clear and useful deprecation warning (or, perhaps even an error simply asking users to switch to the JAX API -- since as far as I can tell the function signature is identical)