JAX FA Benchmarking Script#351

Micky774 wants to merge 10 commits intodevROCm/TransformerEngine:devfrom

zain/jax-benchROCm/TransformerEngine:zain/jax-benchCopy head branch name to clipboard