I Am not able to import trax library in jupyter notebook

Question

import os
import nltk
import trax
from trax import layers as tl
from trax.fastmath import numpy as fastnp

Error:

ValueError                                Traceback (most recent call last)
Cell In[1], line 3
      1 import os
      2 import nltk
----> 3 import trax
      4 from trax import layers as tl
      5 from trax.fastmath import numpy as fastnp

File ~\AppData\Roaming\Python\Python311\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 ~\AppData\Roaming\Python\Python311\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 ~\AppData\Roaming\Python\Python311\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>

I installed trax (pip install trax) and tensorflow as usual in jupyter. but it gives me this error. Currently I am using python 3.11 version.

Any suggestion might help.