trax ValueError: no signature found for builtin function

Description

I'm using Tensorflow on my M1 Mac, but importing Trax throws this error. I've tried envs with Tensorflow versions 2.5.0, 2.6.0, and 2.7.0.

Environment information

OS: MacOS 12.2
Arch: ARM64

$ pip freeze | grep trax
trax==1.2.4

$ pip freeze | grep tensor
mesh-tensorflow==0.1.19
tensor2tensor==1.15.5
tensorboard @ file:///home/conda/feedstock_root/build_artifacts/tensorboard_1629677129676/work/tensorboard-2.6.0-py3-none-any.whl
tensorboard-data-server @ file:///Users/runner/miniforge3/conda-bld/tensorboard-data-server_1636045914528/work/tensorboard_data_server-0.6.0-py3-none-macosx_11_0_arm64.whl
tensorboard-plugin-wit @ file:///home/conda/feedstock_root/build_artifacts/tensorboard-plugin-wit_1641458951060/work/tensorboard_plugin_wit-1.8.1-py3-none-any.whl
tensorflow-datasets==4.5.2
tensorflow-estimator==2.8.0
tensorflow-gan==2.1.0
tensorflow-hub==0.12.0
tensorflow-macos==2.6.0
tensorflow-metadata==1.6.0
tensorflow-metal==0.2.0
tensorflow-probability==0.7.0

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

$ python -V
Python 3.9.10

For bugs: reproduction and error logs

# Steps to reproduce:
import trax

# Error logs:

>>> import trax
/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/jax/_src/lib/__init__.py:33: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/__init__.py", line 18, in <module>
    from trax import layers
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/layers/__init__.py", line 23, in <module>
    from trax.layers.activation_fns import *
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/layers/activation_fns.py", line 26, in <module>
    from trax import math
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/math/__init__.py", line 20, in <module>
    from trax.math import tf as tf_math
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/math/tf.py", line 26, in <module>
    from trax.tf_numpy import extensions as tf_np_extensions
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/extensions/__init__.py", line 23, in <module>
    from trax.tf_numpy.extensions.extensions import *
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/extensions/extensions.py", line 30, in <module>
    from trax.tf_numpy.numpy import arrays
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/numpy/__init__.py", line 27, in <module>
    from trax.tf_numpy.numpy.array_creation import *
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/numpy/array_creation.py", line 295, in <module>
    @utils.np_doc(np.asarray)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/numpy/utils.py", line 192, in np_doc
    np_sig = _np_signature(np_fun)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/trax/tf_numpy/numpy/utils.py", line 136, in _np_signature
    return funcsigs.signature(f)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/tensorflow-2.6.0/lib/python3.9/site-packages/funcsigs/__init__.py", line 185, in signature
    raise ValueError(msg)
ValueError: no signature found for builtin function <built-in function asarray>

Feb 11 '22 23:02 alistairmayo

Same environment, same issue.

% pip freeze | grep trax
trax==1.2.4

% pip freeze | grep tensor
mesh-tensorflow==0.1.19
tensor2tensor==1.15.7
tensorboard==2.8.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow-addons==0.16.1
tensorflow-datasets==4.5.2
tensorflow-gan==2.1.0
tensorflow-hub==0.12.0
tensorflow-macos==2.8.0
tensorflow-metadata==1.7.0
tensorflow-metal==0.4.0
tensorflow-probability==0.7.0

% python -V
Python 3.9.10

Whole environment:

absl-py==1.0.0
astunparse==1.6.3
bz2file==0.98
cached-property @ file:///home/conda/feedstock_root/build_artifacts/cached_property_1615209429212/work
cachetools==5.0.0
certifi==2021.10.8
charset-normalizer==2.0.12
click==8.0.4
cloudpickle==2.0.0
decorator==5.1.1
dill==0.3.4
dopamine-rl==3.0.1
Flask==2.0.3
flatbuffers==2.0
funcsigs==1.0.2
future==0.18.2
gast==0.5.3
gevent==21.12.0
gin-config==0.5.0
google-api-core==2.7.1
google-api-python-client==2.40.0
google-auth==2.6.0
google-auth-httplib2==0.1.0
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
googleapis-common-protos==1.55.0
greenlet==1.1.2
grpcio @ file:///Users/runner/miniforge3/conda-bld/grpcio_1645230369499/work
gunicorn==20.1.0
gym==0.23.0
gym-notices==0.0.6
h5py @ file:///Users/runner/miniforge3/conda-bld/h5py_1637964045648/work
httplib2==0.20.4
idna==3.3
importlib-metadata==4.11.2
itsdangerous==2.1.1
jax==0.3.1
jaxlib==0.3.0
Jinja2==3.0.3
keras==2.8.0
Keras-Preprocessing==1.1.2
kfac==0.2.0
libclang==13.0.0
Markdown==3.3.6
MarkupSafe==2.1.0
mesh-tensorflow==0.1.19
mpmath==1.2.1
numpy @ file:///Users/runner/miniforge3/conda-bld/numpy_1640083125610/work
oauth2client==4.1.3
oauthlib==3.2.0
opencv-python==4.5.5.64
opt-einsum==3.3.0
Pillow==9.0.1
promise==2.3
protobuf==3.19.4
pyasn1==0.4.8
pyasn1-modules==0.2.8
pyparsing==3.0.7
pypng==0.0.21
requests==2.27.1
requests-oauthlib==1.3.1
rsa==4.8
scipy==1.8.0
six==1.15.0
sympy==1.10
tensor2tensor==1.15.7
tensorboard==2.8.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow-addons==0.16.1
tensorflow-datasets==4.5.2
tensorflow-gan==2.1.0
tensorflow-hub==0.12.0
tensorflow-macos==2.8.0
tensorflow-metadata==1.7.0
tensorflow-metal==0.4.0
tensorflow-probability==0.7.0
termcolor==1.1.0
tf-estimator-nightly==2.8.0.dev2021122109
tf-slim==1.1.0
tqdm==4.63.0
trax==1.2.4
typeguard==2.13.3
typing_extensions==4.1.1
uritemplate==4.1.1
urllib3==1.26.8
Werkzeug==2.0.3
wrapt==1.14.0
zipp==3.7.0
zope.event==4.5.0
zope.interface==5.4.0

Clean install with conda env.