Enable mixed scan strategy for models with Engram

[RissyRan]

Description

Refactor and enable scan for DeepSeek custom model with Engram layers, config link

Due to challenges with random layer IDs across the entire scan block, the model is currently partitioned segments in order. For instance, we have 10 layers, and first 5 as Dense, and last 5 as MoE, and Engram is at 3rd, and 9th.

* Scanned - Dense layers 0-1
* Unscanned - Engram layer 2
* Scanned - Dense layers 3-4
* Scanned - MoE layers 5-7
* Unscanned - Engram layer 8
* Scanned - MoE layers 9

Tests

• Added unit test - deepseek_scan_engram_test.py, and expect to pass
• Un-scan vs. scan (deepseek-custom, max_target_length=4096 per_device_batch_size=16 on v5p) link - Obvious performance gains and reduced memory footprint.
• No impact for existing deepseek related models - DeepSeek v2 scan - before vs. after change - no change link

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 92.30769% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/decoders.py	92.30%	1 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

[gagika]

the logic looks good to me.
could you try to do through a helper function and some code re-use?

[gobbleturk]

Thanks Ran!