Unable to import trax

[JoseJamesMD]

Description

I installed Trax but cannot import

Environment information

OS: Ubuntu 20.04.5 LTS

$ pip freeze | grep trax
trax==1.4.1

$ pip freeze | grep tensor
tensorboard==2.11.2
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.11.0
tensorflow-datasets==4.8.2
tensorflow-estimator==2.11.0
tensorflow-hub==0.12.0
tensorflow-io-gcs-filesystem==0.31.0
tensorflow-metadata==1.12.0
tensorflow-text==2.11.0

$ pip freeze | grep jax
jax==0.4.4
jaxlib==0.4.4

$ python -V
Python 3.9.16

### For bugs: reproduction and error logs

# Steps to reproduce:

I installed trax via pip. Did not notice any errors.

$ pip install -U trax
Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: trax in ./.local/lib/python3.8/site-packages (1.4.1)
Requirement already satisfied: scipy in ./.local/lib/python3.8/site-packages (from trax) (1.10.1)
Requirement already satisfied: gym in ./.local/lib/python3.8/site-packages (from trax) (0.26.2)
Requirement already satisfied: numpy in ./.local/lib/python3.8/site-packages (from trax) (1.24.2)
Requirement already satisfied: funcsigs in ./.local/lib/python3.8/site-packages (from trax) (1.0.2)
Requirement already satisfied: tensorflow-text in ./.local/lib/python3.8/site-packages (from trax) (2.11.0)
Requirement already satisfied: jax in ./.local/lib/python3.8/site-packages (from trax) (0.4.4)
Requirement already satisfied: gin-config in ./.local/lib/python3.8/site-packages (from trax) (0.5.0)
Requirement already satisfied: psutil in /usr/lib/python3/dist-packages (from trax) (5.5.1)
Requirement already satisfied: six in /usr/lib/python3/dist-packages (from trax) (1.14.0)
Requirement already satisfied: jaxlib in ./.local/lib/python3.8/site-packages (from trax) (0.4.4)
Requirement already satisfied: matplotlib in /usr/lib/python3/dist-packages (from trax) (3.1.2)
Requirement already satisfied: absl-py in ./.local/lib/python3.8/site-packages (from trax) (1.4.0)
Requirement already satisfied: tensorflow-datasets in ./.local/lib/python3.8/site-packages (from trax) (4.8.2)
Requirement already satisfied: gym-notices>=0.0.4 in ./.local/lib/python3.8/site-packages (from gym->trax) (0.0.8)
Requirement already satisfied: cloudpickle>=1.2.0 in /usr/lib/python3/dist-packages (from gym->trax) (1.3.0)
Requirement already satisfied: importlib-metadata>=4.8.0 in ./.local/lib/python3.8/site-packages (from gym->trax) (6.0.0)
Requirement already satisfied: opt-einsum in ./.local/lib/python3.8/site-packages (from jax->trax) (3.3.0)
Requirement already satisfied: protobuf>=3.12.2 in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (3.19.6)
Requirement already satisfied: termcolor in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (2.2.0)
Requirement already satisfied: wrapt in /usr/lib/python3/dist-packages (from tensorflow-datasets->trax) (1.11.2)
Requirement already satisfied: promise in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (2.3)
Requirement already satisfied: importlib-resources in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (5.12.0)
Requirement already satisfied: dm-tree in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (0.1.8)
Requirement already satisfied: toml in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (0.10.2)
Requirement already satisfied: dill in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (0.3.6)
Requirement already satisfied: requests>=2.19.0 in /usr/lib/python3/dist-packages (from tensorflow-datasets->trax) (2.22.0)
Requirement already satisfied: tensorflow-metadata in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (1.12.0)
Requirement already satisfied: click in /usr/lib/python3/dist-packages (from tensorflow-datasets->trax) (7.0)
Requirement already satisfied: etils[enp,epath]>=0.9.0 in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (1.0.0)
Requirement already satisfied: tqdm in ./.local/lib/python3.8/site-packages (from tensorflow-datasets->trax) (4.64.1)
Requirement already satisfied: tensorflow-hub>=0.8.0 in ./.local/lib/python3.8/site-packages (from tensorflow-text->trax) (0.12.0)
Requirement already satisfied: tensorflow<2.12,>=2.11.0 in ./.local/lib/python3.8/site-packages (from tensorflow-text->trax) (2.11.0)
Requirement already satisfied: zipp in ./.local/lib/python3.8/site-packages (from etils[enp,epath]>=0.9.0->tensorflow-datasets->trax) (3.15.0)
Requirement already satisfied: typing_extensions in ./.local/lib/python3.8/site-packages (from etils[enp,epath]>=0.9.0->tensorflow-datasets->trax) (4.5.0)
Requirement already satisfied: setuptools in /usr/lib/python3/dist-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (45.2.0)
Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.31.0)
Requirement already satisfied: astunparse>=1.6.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (1.6.3)
Requirement already satisfied: tensorflow-estimator<2.12,>=2.11.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.11.0)
Requirement already satisfied: gast<=0.4.0,>=0.2.1 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.4.0)
Requirement already satisfied: keras<2.12,>=2.11.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.11.0)
Requirement already satisfied: libclang>=13.0.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (15.0.6.1)
Requirement already satisfied: tensorboard<2.12,>=2.11 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.11.2)
Requirement already satisfied: google-pasta>=0.1.1 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.2.0)
Requirement already satisfied: grpcio<2.0,>=1.24.3 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (1.51.3)
Requirement already satisfied: packaging in /usr/lib/python3/dist-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (20.3)
Requirement already satisfied: h5py>=2.9.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (3.8.0)
Requirement already satisfied: flatbuffers>=2.0 in ./.local/lib/python3.8/site-packages (from tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (23.1.21)
Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in ./.local/lib/python3.8/site-packages (from tensorflow-metadata->tensorflow-datasets->trax) (1.58.0)
Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/lib/python3/dist-packages (from astunparse>=1.6.0->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.34.2)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.6.1)
Requirement already satisfied: google-auth<3,>=1.6.3 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.16.1)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (1.8.1)
Requirement already satisfied: werkzeug>=1.0.1 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.2.3)
Requirement already satisfied: markdown>=2.6.8 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (3.4.1)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in ./.local/lib/python3.8/site-packages (from tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.4.6)
Requirement already satisfied: pyasn1-modules>=0.2.1 in ./.local/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.2.8)
Requirement already satisfied: rsa<5,>=3.1.4 in ./.local/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (4.9)
Requirement already satisfied: cachetools<6.0,>=2.0.0 in ./.local/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (5.3.0)
Requirement already satisfied: requests-oauthlib>=0.7.0 in ./.local/lib/python3.8/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (1.3.1)
Requirement already satisfied: MarkupSafe>=2.1.1 in ./.local/lib/python3.8/site-packages (from werkzeug>=1.0.1->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (2.1.2)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in ./.local/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (0.4.8)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow<2.12,>=2.11.0->tensorflow-text->trax) (3.1.0)



# Error logs:
unable to import trax

$ python
Python 3.9.16 (main, Jan 11 2023, 16:05:54)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import trax
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'trax'

[JoseJamesMD]

nevermind, issue with PATH for pip which did not install to my conda environment