google / trax

Trax — Deep Learning with Clear Code and Speed
Apache License 2.0
8.01k stars 813 forks source link

Is possible Linformer algorithm ? #1780

Open andr1972 opened 11 months ago

andr1972 commented 11 months ago

Description

I'm starting to test Trax. Is possible Linformer algorithm in future? In Linformer paper I read that it is faster than Reformer.

Environment information

OS: Linux Mint 21.1, https://colab.research.google.com

$ pip freeze | grep trax (none yet)

$ pip freeze | grep tensor (colab) pytensor==2.10.1 tensorboard==2.12.3 tensorboard-data-server==0.7.1 tensorflow==2.12.0 tensorflow-datasets==4.9.2 tensorflow-estimator==2.12.0 tensorflow-gcs-config==2.12.0 tensorflow-hub==0.13.0 tensorflow-io-gcs-filesystem==0.32.0 tensorflow-metadata==1.13.1 tensorflow-probability==0.20.1 tensorstore==0.1.40

$ pip freeze | grep jax jax==0.4.13 jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.13+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl#sha256=af30095a0adf342b837a0ed9607e13177ee66f4e654c031a383aa546cd21d815