Reject KV head counts that cannot be sharded across the tensor-parallel mesh. - #4811

Reject KV head counts that cannot be sharded across the tensor-parallel mesh.#4811

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

mainAI-Hypercomputer/maxtext:mainfrom

test_961244739AI-Hypercomputer/maxtext:test_961244739Copy head branch name to clipboard