Trax Throw error while installing on MAC book pro with M1 chip

Description

I created and environment using miniconda supported with M1 chip, after that we installed Tensorflow2.4 which can use GPU and tried installing Trax but it throws following error.

Collecting trax Using cached trax-1.3.7-py2.py3-none-any.whl (521 kB) Collecting funcsigs Using cached funcsigs-1.0.2-py2.py3-none-any.whl (17 kB) Collecting tensorflow-datasets Using cached tensorflow_datasets-4.2.0-py3-none-any.whl (3.7 MB) Requirement already satisfied: six in ./miniforge3/envs/tf24_env/lib/python3.8/site-packages (from trax) (1.15.0) Requirement already satisfied: numpy in ./miniforge3/envs/tf24_env/lib/python3.8/site-packages (from trax) (1.18.5) Collecting trax Using cached trax-1.3.6-py2.py3-none-any.whl (468 kB) Using cached trax-1.3.5-py2.py3-none-any.whl (416 kB) Using cached trax-1.3.4-py2.py3-none-any.whl (366 kB) Using cached trax-1.3.3-py2.py3-none-any.whl (365 kB) Using cached trax-1.3.2-py2.py3-none-any.whl (365 kB) Using cached trax-1.3.1-py2.py3-none-any.whl (347 kB) Using cached trax-1.3-py2.py3-none-any.whl (346 kB) Using cached trax-1.2.4-py2.py3-none-any.whl (426 kB) Collecting gin-config Using cached gin_config-0.4.0-py2.py3-none-any.whl (46 kB) Collecting trax Using cached trax-1.2.3-py2.py3-none-any.whl (371 kB) Using cached trax-1.2.2-py2.py3-none-any.whl (311 kB) Using cached trax-1.2.1-py2.py3-none-any.whl (311 kB) Using cached trax-1.2.0-py2.py3-none-any.whl (311 kB) Using cached trax-1.1.4-py2.py3-none-any.whl (309 kB) Using cached trax-1.1.3-py2.py3-none-any.whl (309 kB) Using cached trax-1.1.2-py2.py3-none-any.whl (283 kB) Using cached trax-1.1.1-py2.py3-none-any.whl (270 kB) Using cached trax-1.1.0-py2.py3-none-any.whl (231 kB) Using cached trax-1.0.1-py2.py3-none-any.whl (211 kB) Using cached trax-1.0.0-py2.py3-none-any.whl (207 kB) ERROR: Cannot install trax==1.0.0, trax==1.0.1, trax==1.1.0, trax==1.1.1, trax==1.1.2, trax==1.1.3, trax==1.1.4, trax==1.2.0, trax==1.2.1, trax==1.2.2, trax==1.2.3, trax==1.2.4, trax==1.3, trax==1.3.1, trax==1.3.2, trax==1.3.3, trax==1.3.4, trax==1.3.5, trax==1.3.6 and trax==1.3.7 because these package versions have conflicting dependencies.

The conflict is caused by: trax 1.3.7 depends on tensorflow-text trax 1.3.6 depends on tensorflow-text trax 1.3.5 depends on tensorflow-text trax 1.3.4 depends on tensorflow-text trax 1.3.3 depends on tensorflow-text trax 1.3.2 depends on tensorflow-text trax 1.3.1 depends on tensorflow-text trax 1.3 depends on tensorflow-text trax 1.2.4 depends on jaxlib trax 1.2.3 depends on jaxlib trax 1.2.2 depends on jaxlib trax 1.2.1 depends on jaxlib trax 1.2.0 depends on jaxlib trax 1.1.4 depends on jaxlib trax 1.1.3 depends on jaxlib trax 1.1.2 depends on jaxlib trax 1.1.1 depends on jaxlib trax 1.1.0 depends on jaxlib trax 1.0.1 depends on jaxlib trax 1.0.0 depends on jaxlib

To fix this you could try to:

1. loosen the range of package versions you've specified
2. remove package versions to allow pip attempt to solve the dependency conflict

ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/user_guide/#fixing-conflicting-dependencies

Environment information

OS: Mac OS Big Sur 11.2.1

$ pip freeze | grep trax
no output

$ pip freeze | grep tensor
# your output here

grpcio @ file:///Users/pratik/Downloads/tensorflow_macos/arm64/grpcio-1.33.2-cp38-cp38-macosx_11_0_arm64.whl
h5py @ file:///Users/pratik/Downloads/tensorflow_macos/arm64/h5py-2.10.0-cp38-cp38-macosx_11_0_arm64.whl
numpy @ file:///Users/pratik/Downloads/tensorflow_macos/arm64/numpy-1.18.5-cp38-cp38-macosx_11_0_arm64.whl
tensorboard==2.4.1
tensorboard-plugin-wit==1.8.0
tensorflow-addons-macos @ file:///Users/pratik/Downloads/tensorflow_macos/arm64/tensorflow_addons_macos-0.1a2-cp38-cp38-macosx_11_0_arm64.whl
tensorflow-estimator==2.4.0
tensorflow-macos @ file:///Users/pratik/Downloads/tensorflow_macos/arm64/tensorflow_macos-0.1a2-cp38-cp38-macosx_11_0_arm64.whl

$ pip freeze | grep jax
# your output here
no output
$ python -V
# your output here
Python 3.8.6

### For bugs: reproduction and error logs
pip install Trax

Steps to reproduce:

...

Error logs:

...

ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/user_guide/#fixing-conflicting-dependencies