jax
flax>=0.7.1
transformer_engine_rocm7==2.15.0.dev0+devrocm7.14.0.dev0.22b25e93720109d1f2e3b475df7c89cfd29aad98.71eb59d
