fix GemmaBackbone.get_layout_map + test

[martin-gorner]

This PR fixes several issues:

1. The default Gemma layout map was still missing the fixes from issue 19496.
2. The PR also adds a test for Gemma + get_layout_map + LoRA. This test fails without the get_layout_map fix in 1).
3. The PR also adds a layout map for ffw_gating_2, which was apparently forgotten (but there was a test for it!). The value chosen has been tested to be the one offering the fastest training on TPU_v3_8.
4. The PR also fixes a typo in the tests "ffw_linearl" => "ffw_linear"
5. Finally, the PR changes the number of heads in the test Gemma config from 4 to 8, otherwise the tests won't pass on TPU

Both tests test_distribution and test_distribution_with_lora now pass on TPU_v3_8.
However, I had to manually change device detection from 'CPU' to 'TPU' in file gemma_backbone_test.py, or the distribution tests would be skipped. I suspect these tests are NOT run nightly and I kindly ask a team member to check this.

[mattdangerw]

Looks good! Just minor comments.