From bf8b3112572e05babddddc6509d6f2e9415a44e6 Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Fri, 9 Oct 2026 01:20:48 -0400 Subject: [PATCH 1/3] Add a torch.compile graph-break regression test for DiNTS Follow-up to #9145. The topology-cache tests added there pin the snapshot mechanism -- that `_node_a_py` / `_arch_code_a_py` agree with their source tensors and that eager still honours mutations -- but nothing asserts the outcome #9144 is actually about: that a deployed DiNTS traces into a single graph. A refactor that branched on the tensors again would keep those tests green while reintroducing all 13 breaks. Assert `torch._dynamo.explain` reports (1 graph, 0 breaks) on the grid from the issue. The same grid reports (14, 13) with `dints.py` from a7904ae1^. Signed-off-by: Patel, Nilaykumar K --- tests/networks/nets/test_dints_network.py | 28 +++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/networks/nets/test_dints_network.py b/tests/networks/nets/test_dints_network.py index fad32c0350..3a6082495c 100644 --- a/tests/networks/nets/test_dints_network.py +++ b/tests/networks/nets/test_dints_network.py @@ -13,10 +13,12 @@ import os import tempfile +import logging import unittest import numpy as np import torch +import torch._dynamo from parameterized import parameterized from monai.networks.nets import DiNTS, TopologyInstance, TopologySearch @@ -338,6 +340,32 @@ def test_eager_honours_caller_owned_arch_code_array(self): self.assertFalse(bool(grid.arch_code_a[0, 0])) +@unittest.skipUnless(torch._dynamo.is_dynamo_supported(), "dynamo is unsupported on this platform") +class TestDintsGraphBreaks(unittest.TestCase): + def test_compiles_to_a_single_graph(self): + """A deployed DiNTS traces into one graph, with no break on a topology branch.""" + # Before https://github.com/Project-MONAI/MONAI/issues/9144 was fixed this grid reported + # 14 graphs / 13 breaks: one break per `node_a` / `arch_code_a` element read. + torch._dynamo.reset() + # dynamo logs thousands of DEBUG records while tracing; pytest replays them all on + # failure, burying the assertion message. + trace_logger = logging.getLogger("torch.__trace") + self.addCleanup(trace_logger.setLevel, trace_logger.level) + trace_logger.setLevel(logging.WARNING) + grid = TopologyInstance( + channel_mul=0.2, num_blocks=6, num_depths=3, use_downsample=True, spatial_dims=3, device="cpu" + ) + net = DiNTS(dints_space=grid, in_channels=1, num_classes=2, spatial_dims=3, use_downsample=True).eval() + + explanation = torch._dynamo.explain(net)(torch.randn(1, 1, 32, 32, 32)) + + self.assertEqual( + (explanation.graph_count, explanation.graph_break_count), + (1, 0), + "DiNTS no longer compiles to a single graph; a topology branch is reading a tensor again (#9144).", + ) + + class TestDintsTS(unittest.TestCase): @parameterized.expand(TEST_CASES_3D + TEST_CASES_2D) def test_script(self, dints_grid_params, dints_params, input_shape, _): From 0e1d2f5ba3651c50856117d658332dcfb000ab90 Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Wed, 21 Oct 2026 01:53:43 -0400 Subject: [PATCH 2/3] Address code rabbit comment for edge case of boundary graph break. Signed-off-by: Patel, Nilaykumar K --- tests/networks/nets/test_dints_network.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/networks/nets/test_dints_network.py b/tests/networks/nets/test_dints_network.py index 3a6082495c..d74721e6c3 100644 --- a/tests/networks/nets/test_dints_network.py +++ b/tests/networks/nets/test_dints_network.py @@ -357,7 +357,9 @@ def test_compiles_to_a_single_graph(self): ) net = DiNTS(dints_space=grid, in_channels=1, num_classes=2, spatial_dims=3, use_downsample=True).eval() - explanation = torch._dynamo.explain(net)(torch.randn(1, 1, 32, 32, 32)) + input_tensor = torch.randn(1, 1, 32, 32, 32) + torch.compile(net, backend="eager", fullgraph=True)(input_tensor) + explanation = torch._dynamo.explain(net)(input_tensor) self.assertEqual( (explanation.graph_count, explanation.graph_break_count), From 60273372d41aa1d4af7d221bf8b75fa129ac1bfa Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Wed, 21 Oct 2026 03:07:11 -0400 Subject: [PATCH 3/3] fix: sort stdlib imports alphabetically to pass isort check `import logging` was placed after `import os` and `import tempfile`, violating the alphabetical order required by MONAI's isort config (profile=black). Move it to the correct position. Signed-off-by: Patel, Nilaykumar K --- tests/networks/nets/test_dints_network.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/networks/nets/test_dints_network.py b/tests/networks/nets/test_dints_network.py index d74721e6c3..902a3d1e81 100644 --- a/tests/networks/nets/test_dints_network.py +++ b/tests/networks/nets/test_dints_network.py @@ -11,9 +11,9 @@ from __future__ import annotations +import logging import os import tempfile -import logging import unittest import numpy as np