Add NaViT: Native Resolution Vision Transformer with Patch n' Pack by vikashg · Pull Request #9011 · Project-MONAI/MONAI

vikashg · 2026-07-24T06:47:36Z

Adds NaViT (monai.networks.nets.NaViT), a Vision Transformer that removes the fixed-resolution constraint of standard ViT by packing multiple variable-size images into a single sequence per batch element.

Key features:

Patch n' Pack: multiple images concatenated into one sequence per group, with a per-image attention mask preventing cross-image attention
Factorised positional embeddings: separate learnable tables per spatial axis, allowing generalisation to unseen resolutions
Token dropout: configurable fraction of patch tokens dropped during training (float or callable)
Attention pooling: learned query attends over each image's tokens to produce a fixed-size per-image representation
QK normalisation: RMS normalisation on queries and keys (ViT-22B style)
2D and 3D support: works for (C, H, W) and (C, H, W, D) inputs

Changes:

monai/networks/nets/navit.py: new NaViT implementation
monai/networks/nets/init.py: export NaViT
tests/networks/nets/test_navit.py: 24 unit tests covering shape, variable resolutions, token dropout, auto-grouping, gradient flow, ill arguments, and forward validation
docs/source/networks.rst: autoclass entry
docs/source/whatsnew_1_5_2.md: feature description
CHANGELOG.md: entry under Unreleased

Fixes # .

Description

A few sentences describing the changes proposed in this pull request.

Types of changes

Non-breaking change (fix or new feature that would not break existing functionality).
Breaking change (fix or new feature that would cause existing functionality to change).
New tests added to cover the changes.
Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
In-line docstrings updated.
Documentation updated, tested make html command in the docs/ folder.

coderabbitai · 2026-07-24T06:47:54Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds the NaViT native-resolution vision transformer for variable-resolution 2D and 3D inputs. The implementation supports patch extraction, token dropout, sequence grouping, packed attention masks, factorized positional embeddings, transformer blocks, attention pooling, and classification. NaViT is exported from monai.networks.nets, documented in Sphinx, listed in the changelog, and covered by validation and behavior tests.

Estimated code review effort: 4 (Complex) | ~60 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: adding the NaViT native-resolution vision transformer with Patch n' Pack.
Description check	✅ Passed	The description explains the implementation, tests, documentation, and change type, but retains duplicated template text and an incomplete issue reference.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 4

🧹 Nitpick comments (1)

monai/networks/nets/navit.py (1)
427-446: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

forward() docstring omits the Raises section.

forward() raises ValueError for wrong ndim (line 475), wrong channel count (line 480), and non-divisible spatial dims (line 484), but none of this is documented. As per path instructions, docstrings should describe raised exceptions in the appropriate Google-style section.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/networks/nets/navit.py` around lines 427 - 446, The forward() docstring
in NaViT should document its ValueError cases. Add a Google-style Raises section
describing ValueError for invalid tensor dimensionality, mismatched channel
count, or spatial dimensions not divisible by patch_size.
Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@monai/networks/nets/navit.py`:
- Around line 344-348: Validate each dimension of image_size in the constructor
before building pos_embed_axes, requiring exact divisibility by patch_size and
raising a clear error otherwise. Use image_size_t and preserve the existing
positional-embedding table sizing only for valid dimensions.
- Around line 314-321: Update the constructor validation near the
hidden_size/num_heads divisibility check to explicitly reject non-positive
num_heads with ValueError before evaluating hidden_size % num_heads. Preserve
the existing ValueError for non-divisible positive values and the documented
validation behavior.
- Around line 492-497: Gate the token-dropout block in the NaViT forward path on
both self.calc_token_dropout being set and self.training, so dropout is applied
only during training. Keep the existing dropout fraction, index selection, and
sequence/position filtering unchanged when training is active.
- Around line 206-218: Update the forward method to preserve the post-attention
residual in x, pass only its normalized value to self.mlp, and add the MLP
output back to the unnormalized residual. Keep the attention residual and
existing return shape unchanged.

---

Nitpick comments:
In `@monai/networks/nets/navit.py`:
- Around line 427-446: The forward() docstring in NaViT should document its
ValueError cases. Add a Google-style Raises section describing ValueError for
invalid tensor dimensionality, mismatched channel count, or spatial dimensions
not divisible by patch_size.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 91dc08cb-fbc3-4e40-8d78-11336fac95b2

📥 Commits

Reviewing files that changed from the base of the PR and between a3d5160 and aa7a76a.

📒 Files selected for processing (6)

CHANGELOG.md
docs/source/networks.rst
docs/source/whatsnew_1_5_2.md
monai/networks/nets/__init__.py
monai/networks/nets/navit.py
tests/networks/nets/test_navit.py

Adds NaViT (monai.networks.nets.NaViT), a Vision Transformer that removes the fixed-resolution constraint of standard ViT by packing multiple variable-size images into a single sequence per batch element. Key features: - Patch n' Pack: multiple images concatenated into one sequence per group, with a per-image attention mask preventing cross-image attention - Factorised positional embeddings: separate learnable tables per spatial axis, allowing generalisation to unseen resolutions - Token dropout: configurable fraction of patch tokens dropped during training (float or callable) - Attention pooling: learned query attends over each image's tokens to produce a fixed-size per-image representation - QK normalisation: RMS normalisation on queries and keys (ViT-22B style) - 2D and 3D support: works for (C, H, W) and (C, H, W, D) inputs Changes: - monai/networks/nets/navit.py: new NaViT implementation - monai/networks/nets/__init__.py: export NaViT - tests/networks/nets/test_navit.py: 24 unit tests covering shape, variable resolutions, token dropout, auto-grouping, gradient flow, ill arguments, and forward validation - docs/source/networks.rst: autoclass entry - docs/source/whatsnew_1_5_2.md: feature description - CHANGELOG.md: entry under Unreleased Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

NaViT depends on einops (optional dependency), so test_navit must be excluded from the minimal CI runner per CONTRIBUTING.md guidelines. Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

…ling - Remove NaViT section from whatsnew_1_5_2.md: 1.5.2 was a security-only patch release; NaViT belongs in the upcoming release under [Unreleased] - CHANGELOG [Unreleased]: factorised -> factorized, normalisation -> normalization - navit.py docstrings and comments: normalisation -> normalization, generalisation -> generalization, factorised -> factorized (MONAI uses American English per CONTRIBUTING.md) Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

Remediation for bot commit aa7a76a ([pre-commit.ci] auto fixes from pre-commit.com hooks) which removed an unused SABlock import from navit.py. That commit was made by pre-commit-ci[bot] and could not carry a Signed-off-by line. Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

- Fix residual stream bug: use x = self.mlp(self.norm(x)) + x (pre-norm) - Gate token dropout with self.training for deterministic eval - Add num_heads <= 0 ValueError guard - Validate image_size divisibility by patch_size in constructor - Add Raises section to forward() docstring - Refactor tests: DEFAULT_KWARGS, consolidated forward-validation tests Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

coderabbitai

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/networks/nets/test_navit.py`:
- Around line 168-170: Complete the docstring for test_forward_validation by
documenting the image_shape parameter and the ValueError assertion condition,
following the repository’s Google-style format; do not alter the test behavior.
- Around line 184-204: Strengthen test_token_dropout_callable and
test_token_dropout_disabled_in_eval by using a recording callable, asserting it
is invoked in training and not invoked in evaluation. Add a training-mode
fixed-float token-dropout test with identical input under different RNG seeds
and assert the outputs differ, preserving the existing shape and eval
determinism checks.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 8301cac5-e1d5-4442-8c55-002f545792fe

📥 Commits

Reviewing files that changed from the base of the PR and between 85df80d and 1941ea1.

📒 Files selected for processing (2)

monai/networks/nets/navit.py
tests/networks/nets/test_navit.py

🚧 Files skipped from review as they are similar to previous changes (1)

monai/networks/nets/navit.py

Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

…ence - test_token_dropout_callable_invoked_during_training: asserts callable IS invoked during training via a recording wrapper - test_token_dropout_callable_not_invoked_during_eval: asserts callable is NOT invoked during eval mode - test_token_dropout_produces_different_outputs_in_training: same input with different RNG seeds produces different outputs (proves dropout runs) - test_token_dropout_disabled_in_eval: preserved existing determinism check Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

vikashg requested review from KumoLiu, Nic-Ma and ericspod as code owners July 24, 2026 06:47

coderabbitai Bot reviewed Jul 24, 2026

View reviewed changes

Comment thread monai/networks/nets/navit.py

Comment thread monai/networks/nets/navit.py

Comment thread monai/networks/nets/navit.py Outdated

Comment thread monai/networks/nets/navit.py Outdated

vikashg and others added 6 commits July 24, 2026 16:44

[pre-commit.ci] auto fixes from pre-commit.com hooks

2912a6e

for more information, see https://pre-commit.ci

fix mypy no-any-return errors in navit.py

f6a72a9

Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

add test_navit to min_tests exclude list

7e9a086

NaViT depends on einops (optional dependency), so test_navit must be excluded from the minimal CI runner per CONTRIBUTING.md guidelines. Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

vikashg force-pushed the vikash/NaViT branch from 1e9babb to 85df80d Compare July 24, 2026 16:45

vikashg added the enhancement New feature or request label Jul 24, 2026

Merge branch 'dev' into vikash/NaViT

22f14e7

ericspod reviewed Jul 31, 2026

View reviewed changes

Comment thread tests/networks/nets/test_navit.py Outdated

ericspod reviewed Jul 31, 2026

View reviewed changes

Comment thread tests/networks/nets/test_navit.py Outdated

coderabbitai Bot reviewed Aug 7, 2026

View reviewed changes

Comment thread tests/networks/nets/test_navit.py

Comment thread tests/networks/nets/test_navit.py Outdated

vikashg added 2 commits August 7, 2026 07:45

style: apply black formatting to navit.py and test_navit.py

580a214

Signed-off-by: Vikash Gupta <write2vikash@gmail.com>

Navigation Menu

Uh oh!

Add NaViT: Native Resolution Vision Transformer with Patch n' Pack - #9011

Add NaViT: Native Resolution Vision Transformer with Patch n' Pack#9011
vikashg wants to merge 10 commits into
Project-MONAI:devfrom
vikashg:vikash/NaViT

vikashg commented Jul 24, 2026

Uh oh!

coderabbitai Bot commented Jul 24, 2026 •

edited

Loading

Reviews paused

Walkthrough

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

vikashg commented Jul 24, 2026

Description

Types of changes

Uh oh!

coderabbitai Bot commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderabbitai Bot commented Jul 24, 2026 •

edited

Loading