jax
flax>=0.7.1
nvidia-cudnn-frontend>=1.25.0
transformer_engine_cu12==2.18.0
