Updates to JAX flash attention kernel.#4578

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

mainAI-Hypercomputer/maxtext:mainfrom

test_952800018AI-Hypercomputer/maxtext:test_952800018Copy head branch name to clipboard