flax compatability issues

[ParuyrGevorgyan]

The flax version is not specified and due to various changes to keep up with the newer JAX versions, flax is also changing quite a bit. A specific flax version in .toml file would be very helpful.

[amacrutherford]

@bsarkar321 could you take a look?

[alexunderch]

also, with recent colab updates, it's required to have jax at least 0.8.0 for python 3.13, maybe, it's necessity to update the requirements