jax
flax>=0.7.1
transformer_engine_rocm7==2.15.0+rocm7.14.0.e7835ed
