trax
trax copied to clipboard
Is possible Linformer algorithm ?
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