Error trying to import t5x

[osamja]

https://github.com/google-research/t5x/commit/bd1fa1c994ef57cbaf8bc1d182cd08779bf5338b#r141544411

After cloning this repo, I'm getting the following error when trying to import the t5x library that seems related to this change? Let me know if I can provide more info to reproduce this

import t5x

TypeError Traceback (most recent call last)

in <cell line: 1>() ----> 1 import t5x

3 frames

/usr/local/lib/python3.10/dist-packages/t5x/init.py in 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 in 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 in 46 import orbax.checkpoint 47 import seqio ---> 48 from t5x import checkpoints 49 from t5x import models 50 from t5x import optimizers

/usr/local/lib/python3.10/dist-packages/t5x/checkpoints.py in 2070 2071 -> 2072 @ocp.args.register_with_handler( 2073 DatasetCheckpointHandler, for_save=True, for_restore=True 2074 )

TypeError: register_with_handler() got an unexpected keyword argument 'for_restore'

[osamja]

After checking out the commit before this, I can confirm there are no issues importing the t5x library @cpgaffney1

[0xdevalias]

Related:

• https://github.com/google-research/t5x/pull/1547

[Apple-jun]

After checking out the commit before this, I can confirm there are no issues importing the t5x library @cpgaffney1

I have been attempting to use several commits preceding the one recommended, such as commit e4e262b, but I am still encountering the same error. I am unsure if there is an error in my approach and would greatly appreciate your guidance. Here are the steps I have taken:

1. Executed the "Setup Environment" cell.
2. Added and ran the following code: !pip install git+https://github.com/google-research/t5x.git@e4e262b5bbd81e7ca42d83203ade3db0600f5779
3. Added and ran this code: import os; os.kill(os.getpid(), 9)
4. Executed the "Imports and Definitions" cell. Subsequently, I received an error message. I would be immensely grateful for any assistance you could provide.

[garzon]

!pip install --upgrade orbax-checkpoint==0.5.8 works as a workaround for me.

[Apple-jun]

!pip install --upgrade orbax-checkpoint==0.5.8 works as a workaround for me.

Thank you very much for your help! Your method was very effective and solved the problem.