jax
flax>=0.7.1
transformer_engine_rocm10==2.19.0.dev0
