[Distillation] base learn-to-init llama attention for distillation

[vlad-karp]

Description

This PR introduces the base implementation for Learn-to-Init (LTI) attention for distillation (llama only be can easily be generalized).

Relevant details and context:

• Problem being solved: Setting up the foundational LTI components required for effective attention distillation of LLaMA models.
• Implementation details:
  • Adds the core LTI logic in a new learn_to_init_layer.py module.
  • Updates the distillation pipeline (distillation_utils.py and train_distill.py) and decoders to support the new attention layer.
  • 2 LTI implementation for GQA - using bi-linear and global linear map options.
  • Configures the system to get LTI student init-time teacher shapes directly from the configuration.

Tests

• Added a new unit test for LearnToInitDense as well as teacher to student weight injection logic.
  src/maxtext/tests/post_training/unit/learn-to-init_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 24.56140% with 172 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/learn_to_init_layer.py	16.44%	127 Missing ⚠️
...text/trainers/post_train/distillation/lti_utils.py	17.39%	36 Missing and 2 partials ⚠️
.../trainers/post_train/distillation/train_distill.py	70.58%	4 Missing and 1 partial ⚠️
src/maxtext/layers/decoders.py	0.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[gagika]

nit: can it be None?
perhaps change default and data type if it can't.

[gagika]

could you throw an exception if it's not layer-scan mode?

[gagika]

for naming convention could you use "_" instead of "-" e.g. LLAMA2_LTI = "llama2_lti"

[gagika]

could you name the test with "_" insead of "-"? e..g. learn_to_init_test.py

[gagika]

is it logged at every step or once? as it's inside _post_process_train_step