Add DeepSeek Engram layer

[shuningjin]

Description

Background

• Paper: https://arxiv.org/pdf/2601.07372
• Reference Code: https://github.com/deepseek-ai/Engram/blob/main/engram_demo_v1.py

What this PR does

Add Engram layer: engram.py

• NgramHashMapping (non-parametric): CompressedTokenizer + hashing logic, convert "input_id" to "ngram hash_token_id"
  • CompressedTokenizer (non-parametric): convert "input_id" to "compresed_input_id"
• Engram (multi-branch): inputs are "ngram hash_token_id" and "transformer state", MultiHeadEmbedding (lookup embedding using hash id as static memory) + context-aware gating (dot product static memory with contextual state) + ShortConv (temporal smoothing)
  • MultiHeadEmbedding: convert ngram hash_token_id to ngram embedding vector
  • ShortConv (multi-branch): depthwise (mix time steps, not mix channel), causal, short means kernel size is small

Add unit test: tests.unit.engram_vs_reference_test

• for each component, verify the output matches that from reference code

Implementation Notes

Placement of: NgramHashMapping

• NgramHashMapping converts vanilla token-ids to hashed ngram token-ids, which Engram consumes for embedding lookup
• Future: I would like to NgramHashMapping and hash_input_ids generation be put in data input pipeline, which is CPU intensive. Just like how we put tokenizer and input_ids generation in pipeline.
  • This is the code skeleton: https://screenshot.googleplex.com/7YYxr4z7UqvBkpN
  • Also: b/478294696#comment5

Multi-branch

• Engram and ShortConv handles multi-branch input and multi-branch output (if mhc_expansion_rate > 1), using nnx.vmap for independent norm per branch
• Future: to be integrated into multi-branch backbone like mHC.

Tests

unit test against reference

python3 -m pytest -v --pyargs tests.unit.engram_vs_reference_test -rP -s

log: https://paste.googleplex.com/5905570101067776

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

Files with missing lines	Patch %	Lines
src/MaxText/layers/engram.py	0.00%	198 Missing ⚠️

📢 Thoughts on this report? Let us know!

[shuningjin]

@gemini-cli /review

[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 a JAX implementation of the DeepSeek Engram layer, along with comprehensive unit tests that validate its behavior against a PyTorch reference. The code is well-structured and the implementation appears to be correct and thorough. The core logic is sound, and the use of vectorization with nnx.vmap is a good practice for performance.

🔍 General Feedback

• Good Testing: The inclusion of unit tests comparing the JAX implementation to a PyTorch reference is excellent. This provides high confidence in the correctness of the implementation.
• Clear Implementation: The code in engram.py is well-commented and organized, making it easy to follow the logic from the original paper.
• TODOs: I've commented on the TODOs left in the code. Addressing them will improve the clarity and robustness of the implementation.

[github-actions]

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

[RissyRan]

I reviewed the test and CompressedTokenizer. Will continue to review the rest part tomorrow.

[RissyRan]

Thanks for the change! I left some initial comments, and may need to go over multihead embedding and conv parts. It should be quick.

[RissyRan]

LGTM in general! Just a few minor comments. I will have a try to integrate this change with a decoder layer tomorrow and see how it goes.

[aireenmei]

Thanks for the work!