Updates to JAX flash attention kernel. #4559

copybara-service[bot] wants to merge 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

test_951577146AI-Hypercomputer/maxtext:test_951577146Copy head branch name to clipboard