trax import error

[dhrsethi1]

Hello,

When I try and import the trax library like this:

import trax

I get the below error. Any help would be greatly appreciated.

Note: I do not face this error when using Google Collab but only Jupyter Notebook and Jupyter Labs. However, I need to do this on Jupyter notebooks/labs

Thanks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[6], line 1
----> 1 import trax

File ~/anaconda3/lib/python3.11/site-packages/trax/__init__.py:18
      1 # coding=utf-8
      2 # Copyright 2020 The Trax Authors.
      3 #
   (...)
     13 # See the License for the specific language governing permissions and
     14 # limitations under the License.
     16 """Trax top level import."""
---> 18 from trax import layers
     19 from trax import lr_schedules as lr
     20 from trax import math

File ~/anaconda3/lib/python3.11/site-packages/trax/layers/__init__.py:23
     19 import gin
     20 # We create a flat layers.* namespace for uniform calling conventions as we
     21 # upstream changes.
     22 # pylint: disable=wildcard-import
---> 23 from trax.layers.activation_fns import *
     24 from trax.layers.attention import *
     25 from trax.layers.base import *

File ~/anaconda3/lib/python3.11/site-packages/trax/layers/activation_fns.py:26
      1 # coding=utf-8
      2 # Copyright 2020 The Trax Authors.
      3 #
   (...)
     15 
     16 # Lint as: python3
     17 """Layers that compute activation functions.
     18 
     19 An activation layer computes element-wise a nonlinear function of the preceding
   (...)
     23 easier experimentation across different activation functions.
     24 """
---> 26 from trax import math
     27 from trax.layers import base
     28 from trax.math import numpy as np

File ~/anaconda3/lib/python3.11/site-packages/trax/math/__init__.py:20
     18 from trax.math import jax as jax_math
     19 from trax.math import numpy as numpy_math
---> 20 from trax.math import tf as tf_math
     21 from trax.math.backend import *  # pylint: disable=wildcard-import
     22 from trax.math.jax import nested_map

File ~/anaconda3/lib/python3.11/site-packages/trax/math/tf.py:26
     23 import tensorflow.compat.v2 as tf
     25 from trax.shapes import ShapeDtype
---> 26 from trax.tf_numpy import extensions as tf_np_extensions
     27 from trax.tf_numpy import numpy as tf_np
     30 def tf_abstract_eval(f):

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/extensions/__init__.py:23
     20 from __future__ import print_function
     22 # pylint: disable=wildcard-import
---> 23 from trax.tf_numpy.extensions.extensions import *

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/extensions/extensions.py:30
     27 import numpy as np
     28 import tensorflow.compat.v2 as tf
---> 30 from trax.tf_numpy.numpy import arrays
     31 from trax.tf_numpy.numpy import random
     32 from trax.tf_numpy.numpy.array_creation import array

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/numpy/__init__.py:27
     24 from trax.tf_numpy.numpy import random
     26 # pylint: disable=wildcard-import
---> 27 from trax.tf_numpy.numpy.array_creation import *
     28 from trax.tf_numpy.numpy.array_manipulation import *
     29 from trax.tf_numpy.numpy.array_methods import *

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/numpy/array_creation.py:295
    291   result_t = utils.cond(utils.greater(ndmin, ndims), true_fn, lambda: result_t)
    292   return arrays.tensor_to_ndarray(result_t)
--> 295 @utils.np_doc(np.asarray)
    296 def asarray(a, dtype=None):
    297   if dtype:
    298     dtype = utils.result_type(dtype)

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/numpy/utils.py:192, in np_doc(np_fun)
    182 def np_doc(np_fun):
    183   """Attachs numpy docstring to a function.
    184 
    185   Args:
   (...)
    190     decorated function.
    191   """
--> 192   np_sig = _np_signature(np_fun)
    193   def decorator(f):
    194     """The decorator."""

File ~/anaconda3/lib/python3.11/site-packages/trax/tf_numpy/numpy/utils.py:136, in _np_signature(f)
    134 """An enhanced funcsigs.signature that can handle numpy.ufunc."""
    135 if not isinstance(f, np.ufunc):
--> 136   return funcsigs.signature(f)
    137 def names_from_num(prefix, n):
    138   if n <= 0:

File ~/anaconda3/lib/python3.11/site-packages/funcsigs/__init__.py:185, in signature(obj)
    182 if isinstance(obj, types.BuiltinFunctionType):
    183     # Raise a nicer error message for builtins
    184     msg = 'no signature found for builtin function {0!r}'.format(obj)
--> 185     raise ValueError(msg)
    187 raise ValueError('callable {0!r} is not supported by signature'.format(obj))

ValueError: no signature found for builtin function <built-in function asarray>