[BUG] 'airio' has no attribute 'data_sources'

[Orion-Zheng]

I found my t5x raise an error from airio after recent updates.
The error information is

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-1-2159df375230>](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <cell line: 2>()
      1 # Restart session before running
----> 2 from t5x import checkpoints

6 frames
[/usr/local/lib/python3.10/dist-packages/t5x/__init__.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     15 """Import API modules."""
     16 
---> 17 import t5x.adafactor
     18 import t5x.checkpoints
     19 import t5x.decoding

[/usr/local/lib/python3.10/dist-packages/t5x/adafactor.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     63 import jax.numpy as jnp
     64 import numpy as np
---> 65 from t5x import utils
     66 from t5x.optimizers import OptimizerDef
     67 from t5x.optimizers import OptimizerState

[/usr/local/lib/python3.10/dist-packages/t5x/utils.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     32 from absl import flags
     33 from absl import logging
---> 34 import airio
     35 import clu.data
     36 import flax

[/usr/local/lib/python3.10/dist-packages/airio/__init__.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     16 # pylint:disable=wildcard-import,g-bad-import-order,g-importing-member
     17 
---> 18 from airio.dataset_iterators import *
     19 from airio.dataset_providers import *
     20 from airio.data_sources import *

[/usr/local/lib/python3.10/dist-packages/airio/dataset_iterators.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     15 """AirIO-specific dataset iterators."""
     16 
---> 17 from airio.grain import dataset_iterators
     18 
     19 PyGrainDatasetIteratorWrapper = dataset_iterators.PyGrainDatasetIteratorWrapper

[/usr/local/lib/python3.10/dist-packages/airio/grain/__init__.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     17 # pylint:disable=wildcard-import,g-bad-import-order,g-importing-member
     18 
---> 19 from airio.grain.data_sources import *
     20 from airio.grain.dataset_iterators import *
     21 from airio.grain.dataset_providers import *

[/usr/local/lib/python3.10/dist-packages/airio/grain/data_sources.py](https://n60c5jt5kfs-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240108-060129_RC00_596559337#) in <module>
     23 
     24 
---> 25 class ArrayRecordDataSource(airio.data_sources.DataSource):
     26   """Wrapper for grain.ArrayRecordDataSource with multiple splits support."""
     27 

AttributeError: partially initialized module 'airio' has no attribute 'data_sources' (most likely due to a circular import)

To reproduce this error, you can simply run code below in Colab

git clone --branch=main https://github.com/google-research/t5x
python -m pip install ./t5x
# restart session before running
from t5x import checkpoints

[tomlimi]

I have possibly related error when running t5x script:

AttributeError: module 'airio' has no attribute 'PyGrainDatasetIteratorWrapper