Add JAX+Flax example

**Describe the use case example you want to see**
Could you please add an example [here](https://github.com/aws/amazon-sagemaker-examples/blob/main/advanced_functionality/jax_bring_your_own/train_deploy_jax.ipynb) on how to use JAX with [Flax](https://github.com/google/flax) in SageMaker?

**How would this example be used? Please describe.**
Inform users how to use SageMaker with Flax, one of most popular deep learning libraries powered by JAX.

Start with advanced_functionality/jax_bring_your_own/train_deploy_jax.ipynb and review the linked Flax project to understand the existing JAX SageMaker example. Add a Flax usage example in that notebook, with the completed example showing users how to use SageMaker with Flax.

Written by the indexing model from the issue text.