Trax can't use in google colab

the error message shows that jax.ops.index_add is no longer valid.

Could anyone show me how to run Trax in google colab?

Looks like you’re hitting some version issues. Maybe try installing a jax that is pre 0.3.2

The functions jax.ops.index_update , jax.ops.index_add , which were deprecated in 0.2.22, have been removed.…

1 Like

I tried this version portfolio. but it doesn’t work.
the error message when I import trax
ImportError: cannot import name ‘EllipsisType’ from ‘typing’ (/usr/lib/python3.10/

looks like trax 1.3.9 was released May 21, 2021, while Python 3.10 was released October 4, 2021. The Python version aligned with the trax release date would have been more in the range of …

Release. 3.9.2
Date. February 19, 2021


Release. 3.9.4
Date. April 04, 2021

You also might need to use earlier releases of the jax... libraries

which are Python 3.9.x era. (pre-3.10)