DeepSeek MHC feature

[RissyRan]

Description

Initial implementation of DeepSeek MHC feature:

• MHC paper: https://arxiv.org/pdf/2512.24880
• Research implementation from 2 repos; however, I noticed they are not exactly the same as paper. So I took those as reference and implement based on paper, to better fit in MaxText integration.
  • Ref1 - weights init and sinkhorn are slightly different, i.e. sinkhorn used log version
  • Ref2 - has extra configs like num_fracs, num_input_views potentially from HC
  • Move the expansion_rate to 3rd dimension to align with Engram reference (highly likely for V4)

Next: model/decoder layer integration using this deepseek-custom config, and test end-to-end.

Tests

• Unit tests for MHC feature: link
• I didn't find a good way to test implementation details in the unit test (only test shape for core part), but plan to test convergence/loss in the coming model integration, compared to non-MHC workload.

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.

[RissyRan]

@gemini-cli /review

[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]

📋 Review Summary

This pull request introduces the DeepSeek Manifold-Constrained Hyper Connections (mHC) feature. The implementation includes new common types, configuration updates, the core mhc layer, and unit tests. The unit tests provide good coverage for the expand/reduce functions and the sinkhorn algorithm.

🔍 General Feedback

• The overall structure of the mHC implementation is clear, and the separation into common types, configuration, and a dedicated layer is well-organized.
• The unit tests for the utility functions (expand, reduce, sinkhorn) are thorough and cover important properties.
• A critical concern has been identified regarding the dependency of weight matrix initialization on batch and sequence dimensions, which should be addressed for robustness and scalability.

[codecov]

Codecov Report

❌ Patch coverage is 97.56098% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/layers/mhc.py	97.40%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[shuningjin]

Thanks for implementing the feature! There seems to be some ambiguity; I left a few questions. I may need another look.

[shuningjin]

Thanks for the updates! Overall looks good. Left some minor comments.

[shuningjin]

Great work!