Cannot import Trax

[Nuna7]

Description

I am trying to use trax following NLP specialization but couldn't import it. I have search for solution and ask help on ChatGPT but could not solve it.
...

Environment information

OS: <Mac Ventura 13.1>

$ pip freeze | grep trax
# trax==1.2.4

$ pip freeze | grep tensor
# mesh-tensorflow==0.1.21
tensor2tensor==1.15.7
tensorboard @ file:///usr/local/conda/envs/tf210/conda-bld/tensorboard_1670959059865/work/tensorboard-2.10.0-py3-none-any.whl
tensorboard-data-server @ file:///var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_91lm1y0ip6/croot/tensorboard-data-server_1670853587549/work/tensorboard_data_server-0.6.1-py3-none-macosx_11_0_arm64.whl
tensorboard-plugin-wit==1.6.0
tensorflow @ file:///usr/local/conda/envs/tf210/conda-bld/tensorflow-base_1671120030233/work/tensorflow_pkg/tensorflow-2.10.0-cp310-cp310-macosx_12_0_arm64.whl
tensorflow-addons==0.19.0
tensorflow-dataset==0.1.1
tensorflow-datasets==4.8.3
tensorflow-estimator @ file:///usr/local/conda/envs/tf210/conda-bld/tensorflow-estimator_1670959621277/work/tensorflow_estimator-2.10.0-py2.py3-none-any.whl
tensorflow-gan==2.1.0
tensorflow-hub==0.13.0
tensorflow-metadata==1.12.0
tensorflow-probability==0.7.0

$ pip freeze | grep jax
# jax==0.4.8
jaxlib==0.4.7

$ python -V
# Python 3.10.9

For bugs: reproduction and error logs

# Steps to reproduce:
...import trax

# Error logs:
...ValueError: no signature found for builtin function <built-in function asarray>

[antonrosenbloom]

Similar background, the same error

MacOS: 13.3.1

trax==1.2.4

mesh-tensorflow==0.1.21
tensor2tensor==1.15.7
tensorboard==2.13.0
tensorboard-data-server==0.7.0
tensorflow==2.13.0rc0
tensorflow-addons==0.20.0
tensorflow-datasets==4.9.2
tensorflow-estimator==2.13.0rc0
tensorflow-gan==2.1.0
tensorflow-hub==0.13.0
tensorflow-macos==2.13.0rc0
tensorflow-metadata==1.13.1
tensorflow-probability==0.7.0
tensorstore==0.1.37

jax==0.4.10
jaxlib==0.4.10

Python 3.8.10

"import trax" results in "ValueError: no signature found for builtin function <built-in function asarray>"

[tuannguyen90]

I got this error when using cpu environment.
Using a gpu-tensorflow environment works for me.

[boreys]

I face the same issue on Mac OS Ventura with M2 chip

[dhrsethi1]

I am facing the exact same error. This error doesn't happen when I use Google collab, but it is unusably slow when training on large datasets. This forced me back to using Jupyter notebooks. Any idea?

[Israr]

+1 on Python 3.9.6 on Mac OS 14.1.2 (23B92)

[vector-space]

+1 on python 3.11.5, win11, cpu only.

seemed that I soved the problem. trax defined a decorator(in trax\tf_numpy\numpy\utils.py) to attach numpy docstring to its own function. that's the place where this problem is generated. I just added a try-finally in the definition of the decorator to ignore the raised error.