ValueError: no signature found for builtin function

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

# 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>