Installation with Python 3.11 - Grain 0.2.0 outdated URL

Hi,

I tried to install airio using the current setup.py, but the installation fails because the dependency grain==0.2.0 cannot be found or installed.
The repository URL for grain in the dependency list seems outdated or missing (the old URL https://github.com/google-research/grain returns 404).
The current repository for grain appears to be here: https://github.com/google/grain

Would it be possible to update the dependency in setup.py to point to the correct repository or provide guidance on how to install grain?

Currently, this blocks the successful installation of airio and projects that depend on it.