Error: AttributeError: module 'optax' has no attribute 'ScaleByLionState'

[shuntacurosu]

The following error occurs when executing the section "Imports and Definitions".

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-6-152851d66c54>](https://localhost:8080/#) in <cell line: 16>()
     14 import seqio
     15 import t5
---> 16 import t5x
     17 
     18 from mt3 import metrics_utils

5 frames
[/usr/local/lib/python3.9/dist-packages/t5x/optimizers.py](https://localhost:8080/#) in OptaxStatePartitionRules()
    215           lambda state, params_axes: optax.ScaleByBeliefState(  # pytype: disable=wrong-arg-types  # numpy-scalars
    216               count=None, mu=params_axes, nu=params_axes),
--> 217       optax.ScaleByLionState:
    218           lambda state, params_axes: optax.ScaleByLionState(  # pytype: disable=wrong-arg-types  # numpy-scalars
    219               count=None, mu=params_axes),

AttributeError: module 'optax' has no attribute 'ScaleByLionState'

[iansimon]

Should be fixed by https://github.com/google-research/t5x/pull/1248