Fix scale to 1.0 for cudnn_flash_jax#3312

Angelogeb wants to merge 1 commit intoAI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

Angelogeb:cudnn_jax-scale-fixAngelogeb/maxtext:cudnn_jax-scale-fixCopy head branch name to clipboard