Correcting MLP gating and linear einsums. Additionally making changes to the cache shape.#445

copybara-service[bot] wants to merge 1 commit intomaingoogle-deepmind/gemma:mainfrom

test_825721314google-deepmind/gemma:test_825721314Copy head branch name to clipboard