jax
flax>=0.7.1
transformer_engine_rocm7==2.10.0+devrocm7.13.0.dev0.45223d0d4709b5ec27afa0a0efb53faf2d3502a3
