bumping version to python 3.11

[Ruturaj4]

@n-gao Is bumping to python 3.11 absolutely necessary?

We are using ferminet repo on rocm which has folx dependency. However, we can't force python 3.11 (At least not immediately). I am sure there are other repositories that depend on folx may suffer the same issue.

[n-gao]

I will update the CI to cover several Python versions and add support starting from 3.10. However, I will most likely do this at the beginning of October. But if you want to open a PR with the necessary fixes, I am happy to review it.

[Ruturaj4]

Thanks!

[n-gao]

I'll keep this open until folx supports older versions again.

[svandenhaute]

Hi @Ruturaj4 , we're also trying to get ferminet running with ROCm (AMD MI250X on the LUMI supercomputer). We seem to be running into a lot of bugs. Would you mind sharing your installation setup? (python, ROCm driver, jax, jaxlib, kfac_jax versions)?