Add support for higher jax and flax version

[rxng8]

Feature request

Currently, transformer does not support jax version greater than 0.4.13, the libraries has changed and no longer compatible

Motivation

Currently, transformer does not support jax version greater than 0.4.13, the libraries has changed and no longer compatible

Your contribution

N/A

[JackWolfard]

JAX has been deprecated as of #38758. See https://www.linkedin.com/posts/lysandredebut_i-have-bittersweet-news-to-share-yesterday-activity-7338966863403528192-om5p/ for more information from @LysandreJik. I would recommend closing this issue in light of the news.

[Rocketknight1]

Yes, unfortunately! Closing as not planned.