Backward compatibility issue for running Deformable DETR

[sachinraja13]

The requirements file for the deformable DETR project uses jax 0.3.17 and flax 0.5.3.

However, I cannot go back in the version beyond 0.4.26 for jax and 0.8.4 for flax since I want to use METAL accelerator which is only compatible with jax 0.4.26 and above.

Because of this, I'm encountering many issues, specifically with the optimizer. It would be very helpful if I can get some help with porting the code to be compatible with newer versions of jax and flax.

Thanks in advance!