Support Deepstack in qwen3-omni

[aireenmei]

Description

Original work #2729

• rebase and minor fixes
• add unit test TestDeepstackProcess

Tests

qwen3_omni_layer_test.py pass locally
The pylint error is unrelated and will be fixed by #3219

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

Files with missing lines	Patch %	Lines
src/maxtext/layers/decoders.py	9.09%	9 Missing and 1 partial ⚠️
src/maxtext/models/models.py	25.00%	3 Missing ⚠️
src/maxtext/layers/encoders.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @aireenmei, 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 effectively implements the Deepstack visual embedding injection into the Qwen3 Omni model. The integration points in the Decoder and vision_encoder look mostly correct, and the vectorized implementation of _deepstack_process using jnp.cumsum is clever and highly efficient.

🔍 General Feedback

• Efficiency Highlight: The use of boolean masks with cumsum for aligning visual embeddings with sequence positions is an excellent, TPU-friendly pattern that avoids slow dynamic slice operations.
• A logic bug in the scan_layers check could cause the model to silently drop visual embeddings if image_masks are not explicitly provided by the caller but deepstack_visual_embeds are present.
• A minor style improvement can be made by making _deepstack_process a static method, aligning with typical JAX pure-function idioms.

[github-actions]

🤖 Hi @aireenmei, 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 PR successfully integrates Deepstack visual embeddings into the decoder layers for the qwen3-omni models. The implementation effectively extends the vision encoder to output deep features and seamlessly processes them within the transformer decoder architecture.

🔍 General Feedback

• The changes correctly decouple deep feature extraction from the main projection logic while adhering to existing layer configurations.
• The unit tests added are comprehensive and correctly validate the scattering of visual tokens across the sequences based on the bidirectional mask.
• A few minor inline suggestions are provided to handle edge cases relating to implicit type promotions and to prevent potential NaN propagation when computing masked visual embeddings.

[hengtaoguo]

LGTM, thank you Aireen, let's get this merged!